def rechunk_bgen(ds: Dataset, output: Union[(PathType, MutableMapping[(str, bytes)])], *, chunk_length: int=10000, chunk_width: int=1000, compressor: Optional[Any]=zarr.Blosc(cname='zstd', clevel=7, shuffle=2), probability_dtype: Optional[DType]='uint8', max_mem: str='4GB', pack: bool=True, tempdir: Optional[PathType]=None) -> Dataset:
    'Rechunk BGEN dataset as Zarr.\n\n    This function will use the algorithm https://rechunker.readthedocs.io/en/latest/\n    to rechunk certain fields in a provided Dataset for better downstream performance.\n    Depending on the system memory available (and the `max_mem` setting) this\n    rechunking may occur without the need of any intermediate data store. Otherwise,\n    approximately as much disk space is required as was needed to store the original\n    BGEN data. Experiments show that this Zarr representation is ~20% larger even\n    with all available optimizations and fairly aggressive compression (i.e. the\n    default `clevel` 7).\n\n    Note that this function is not evaluated lazily. The rechunking algorithm\n    will run inline so calls to it may be slow. The resulting Dataset is\n    generated based on the final, serialized Zarr data.\n\n    Parameters\n    ----------\n    ds\n        Dataset to rechunk, typically the result from `read_bgen`.\n    output\n        Zarr store or path to directory in file system.\n    chunk_length\n        Length (number of variants) of chunks in which data are stored, by default 10_000.\n    chunk_width\n        Width (number of samples) to use when storing chunks in output, by default 1_000.\n    compressor\n        Zarr compressor, no compression is used when set as None.\n    probability_dtype\n        Data type used to encode genotype probabilities, must be either uint8 or uint16.\n        Setting this parameter results in a loss of precision. If None, probabilities\n        will not be altered when stored.\n    max_mem\n        The amount of memory (in bytes) that workers are allowed to use. A string\n        (e.g. 100MB) can also be used.\n    pack\n        Whether or not to optimize variable representations by removing unnecessary\n        dimensions and elements. This includes storing 2 genotypes instead of 3, omitting\n        dosage and collapsing the genotype probability mask to 2 dimensions. All of\n        the above are restored in the resulting Dataset at the expense of extra\n        computations on read.\n    tempdir\n        Temporary directory where intermediate files are stored. The default None means\n        use the system default temporary directory.\n\n    Warnings\n    --------\n    This functional is only applicable to diploid, bi-allelic BGEN datasets.\n\n    Returns\n    -------\n    Dataset\n        The rechunked dataset.\n    '
    if isinstance(output, Path):
        output = str(output)
    chunk_length = min(chunk_length, ds.dims['variants'])
    chunk_width = min(chunk_width, ds.dims['samples'])
    if pack:
        ds = pack_variables(ds)
    encoding = encode_variables(ds, chunk_length=chunk_length, chunk_width=chunk_width, compressor=compressor, probability_dtype=probability_dtype)
    target_chunks = {var: encoding[var]['chunks'] for var in encoding if ('chunks' in encoding[var])}
    target_options = {var: {k: v for (k, v) in encoding[var].items() if (k != 'chunks')} for var in encoding}
    with tempfile.TemporaryDirectory(prefix='bgen_to_zarr_', suffix='.zarr', dir=tempdir) as tmpdir:
        rechunked = rechunker_api.rechunk(ds, max_mem=max_mem, target_chunks=target_chunks, target_store=output, target_options=target_options, temp_store=tmpdir, executor='dask')
        rechunked.execute()
    zarr.consolidate_metadata(output)
    ds: Dataset = xr.open_zarr(output, concat_characters=False)
    if pack:
        ds = unpack_variables(ds)
    return ds
Rechunk BGEN dataset as Zarr.

This function will use the algorithm https://rechunker.readthedocs.io/en/latest/
to rechunk certain fields in a provided Dataset for better downstream performance.
Depending on the system memory available (and the `max_mem` setting) this
rechunking may occur without the need of any intermediate data store. Otherwise,
approximately as much disk space is required as was needed to store the original
BGEN data. Experiments show that this Zarr representation is ~20% larger even
with all available optimizations and fairly aggressive compression (i.e. the
default `clevel` 7).

Note that this function is not evaluated lazily. The rechunking algorithm
will run inline so calls to it may be slow. The resulting Dataset is
generated based on the final, serialized Zarr data.

Parameters
----------
ds
    Dataset to rechunk, typically the result from `read_bgen`.
output
    Zarr store or path to directory in file system.
chunk_length
    Length (number of variants) of chunks in which data are stored, by default 10_000.
chunk_width
    Width (number of samples) to use when storing chunks in output, by default 1_000.
compressor
    Zarr compressor, no compression is used when set as None.
probability_dtype
    Data type used to encode genotype probabilities, must be either uint8 or uint16.
    Setting this parameter results in a loss of precision. If None, probabilities
    will not be altered when stored.
max_mem
    The amount of memory (in bytes) that workers are allowed to use. A string
    (e.g. 100MB) can also be used.
pack
    Whether or not to optimize variable representations by removing unnecessary
    dimensions and elements. This includes storing 2 genotypes instead of 3, omitting
    dosage and collapsing the genotype probability mask to 2 dimensions. All of
    the above are restored in the resulting Dataset at the expense of extra
    computations on read.
tempdir
    Temporary directory where intermediate files are stored. The default None means
    use the system default temporary directory.

Warnings
--------
This functional is only applicable to diploid, bi-allelic BGEN datasets.

Returns
-------
Dataset
    The rechunked dataset.
def rechunk_bgen(ds: Dataset, output: Union[(PathType, MutableMapping[(str, bytes)])], *, chunk_length: int=10000, chunk_width: int=1000, compressor: Optional[Any]=zarr.Blosc(cname='zstd', clevel=7, shuffle=2), probability_dtype: Optional[DType]='uint8', max_mem: str='4GB', pack: bool=True, tempdir: Optional[PathType]=None) -> Dataset:
    'Rechunk BGEN dataset as Zarr.\n\n    This function will use the algorithm https://rechunker.readthedocs.io/en/latest/\n    to rechunk certain fields in a provided Dataset for better downstream performance.\n    Depending on the system memory available (and the `max_mem` setting) this\n    rechunking may occur without the need of any intermediate data store. Otherwise,\n    approximately as much disk space is required as was needed to store the original\n    BGEN data. Experiments show that this Zarr representation is ~20% larger even\n    with all available optimizations and fairly aggressive compression (i.e. the\n    default `clevel` 7).\n\n    Note that this function is not evaluated lazily. The rechunking algorithm\n    will run inline so calls to it may be slow. The resulting Dataset is\n    generated based on the final, serialized Zarr data.\n\n    Parameters\n    ----------\n    ds\n        Dataset to rechunk, typically the result from `read_bgen`.\n    output\n        Zarr store or path to directory in file system.\n    chunk_length\n        Length (number of variants) of chunks in which data are stored, by default 10_000.\n    chunk_width\n        Width (number of samples) to use when storing chunks in output, by default 1_000.\n    compressor\n        Zarr compressor, no compression is used when set as None.\n    probability_dtype\n        Data type used to encode genotype probabilities, must be either uint8 or uint16.\n        Setting this parameter results in a loss of precision. If None, probabilities\n        will not be altered when stored.\n    max_mem\n        The amount of memory (in bytes) that workers are allowed to use. A string\n        (e.g. 100MB) can also be used.\n    pack\n        Whether or not to optimize variable representations by removing unnecessary\n        dimensions and elements. This includes storing 2 genotypes instead of 3, omitting\n        dosage and collapsing the genotype probability mask to 2 dimensions. All of\n        the above are restored in the resulting Dataset at the expense of extra\n        computations on read.\n    tempdir\n        Temporary directory where intermediate files are stored. The default None means\n        use the system default temporary directory.\n\n    Warnings\n    --------\n    This functional is only applicable to diploid, bi-allelic BGEN datasets.\n\n    Returns\n    -------\n    Dataset\n        The rechunked dataset.\n    '
    if isinstance(output, Path):
        output = str(output)
    chunk_length = min(chunk_length, ds.dims['variants'])
    chunk_width = min(chunk_width, ds.dims['samples'])
    if pack:
        ds = pack_variables(ds)
    encoding = encode_variables(ds, chunk_length=chunk_length, chunk_width=chunk_width, compressor=compressor, probability_dtype=probability_dtype)
    target_chunks = {var: encoding[var]['chunks'] for var in encoding if ('chunks' in encoding[var])}
    target_options = {var: {k: v for (k, v) in encoding[var].items() if (k != 'chunks')} for var in encoding}
    with tempfile.TemporaryDirectory(prefix='bgen_to_zarr_', suffix='.zarr', dir=tempdir) as tmpdir:
        rechunked = rechunker_api.rechunk(ds, max_mem=max_mem, target_chunks=target_chunks, target_store=output, target_options=target_options, temp_store=tmpdir, executor='dask')
        rechunked.execute()
    zarr.consolidate_metadata(output)
    ds: Dataset = xr.open_zarr(output, concat_characters=False)
    if pack:
        ds = unpack_variables(ds)
    return ds