diff --git a/tests/test_create.py b/tests/test_create.py index 13e0fac..3685534 100644 --- a/tests/test_create.py +++ b/tests/test_create.py @@ -61,6 +61,21 @@ def test_create__single(): assert root["variant_allele"][0, 1] == "T" +def test_create__override_samples_chunk_size(): + vcz1 = make_vcz( + [0], + [100], + [["A", "T"]], + sample_id=["S1"], + call_genotype=[[[0, 0]]], + samples_chunk_size=1, + ) + vcz_out = zarr.storage.MemoryStore() + create(vcz_out, vcz1, samples_chunk_size=2) + root = zarr.open(vcz_out) + assert root["call_genotype"].chunks[1] == 2 + + def test_create__no_match(): # Disjoint alts at same position → 2 output variants vcz1 = make_vcz([0], [100], [["A", "T"]]) diff --git a/vczstore/cli.py b/vczstore/cli.py index e15bc70..8dfd047 100644 --- a/vczstore/cli.py +++ b/vczstore/cli.py @@ -98,16 +98,23 @@ def append(vcz1, vcz2, verbose, backend_storage, io_concurrency, require_direct_ @click.command() @click.argument("vcz_out", type=click.Path()) @click.argument("vczs", nargs=-1, type=click.Path()) +@click.option( + "--samples-chunk-size", + type=click.IntRange(min=1), + default=None, + help="Chunk size in the samples dimension", +) @verbose @progress @backend_storage -def create(vcz_out, vczs, verbose, progress, backend_storage): +def create(vcz_out, vczs, samples_chunk_size, verbose, progress, backend_storage): """Create a new, empty store VCZ_OUT using merged variants from VCZS""" setup_logging(verbose) call_or_error( create_function, vcz_out, *vczs, + samples_chunk_size=samples_chunk_size, show_progress=progress, backend_storage=backend_storage, ) diff --git a/vczstore/create.py b/vczstore/create.py index 3559087..11b565c 100644 --- a/vczstore/create.py +++ b/vczstore/create.py @@ -521,7 +521,9 @@ def _compute_merged_variants( ) -def create(vcz_out, *vczs, show_progress=False, backend_storage=None) -> None: +def create( + vcz_out, *vczs, samples_chunk_size=None, show_progress=False, backend_storage=None +) -> None: """Create a new, empty store vcz_out using merged variants from vczs using -m none semantics with stable variant ordering. @@ -555,6 +557,8 @@ def create(vcz_out, *vczs, show_progress=False, backend_storage=None) -> None: root1 = zarr.open(vcz1, mode="r") if len(vczs) == 1: + n_variants = root1["variant_contig"].shape[0] + out_root = open_zarr( vcz_out, mode="w", @@ -713,13 +717,17 @@ def create(vcz_out, *vczs, show_progress=False, backend_storage=None) -> None: if var.startswith("call_"): arr = root1[var] shape = (n_variants, 0) + arr.shape[2:] - # TODO: should allow sample chunk size to be overridden/enforced here + chunks = arr.chunks + if samples_chunk_size is None: + chunks = arr.chunks + else: + chunks = (arr.chunks[0], samples_chunk_size) + arr.chunks[2:] create_empty_group_array( out_root, var, shape=shape, dtype=arr.dtype, - chunks=arr.chunks, + chunks=chunks, compressor=get_compressor_config(arr), dimension_names=array_dims(arr), )