Allow samples chunk size to be set by create command (#107)

Author: tomwhite

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"]])

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,
     )

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),
                 )