Add max_v_chunks option to encode

Closes #57

Author: jeromekelleher

bio2zarr/cli.py

@@ -15,9 +15,21 @@
 )
 
 # TODO help text
-chunk_length = click.option("-l", "--chunk-length", type=int, default=None)
+chunk_length = click.option(
+    "-l",
+    "--chunk-length",
+    type=int,
+    default=None,
+    help="Chunk size in the variants dimension",
+)
 
-chunk_width = click.option("-w", "--chunk-width", type=int, default=None)
+chunk_width = click.option(
+    "-w",
+    "--chunk-width",
+    type=int,
+    default=None,
+    help="Chunk size in the samples dimension",
+)
 
 version = click.version_option(version=provenance.__version__)
 
@@ -83,13 +95,30 @@ def mkschema(if_path):
 @click.argument("if_path", type=click.Path())
 @click.argument("zarr_path", type=click.Path())
 @verbose
-@click.option("-s", "--schema", default=None)
-# TODO: these are mutually exclusive with schema, tell click this
+@click.option("-s", "--schema", default=None, type=click.Path(exists=True))
 @chunk_length
 @chunk_width
+@click.option(
+    "-V",
+    "--max-variant-chunks",
+    type=int,
+    default=None,
+    help=(
+        "Truncate the output in the variants dimension to have "
+        "this number of chunks. Mainly intended to help with "
+        "schema tuning."
+    ),
+)
 @worker_processes
 def encode(
-    if_path, zarr_path, verbose, schema, chunk_length, chunk_width, worker_processes
+    if_path,
+    zarr_path,
+    verbose,
+    schema,
+    chunk_length,
+    chunk_width,
+    max_variant_chunks,
+    worker_processes,
 ):
     """
     Encode intermediate format (see explode) to vcfzarr
@@ -101,6 +130,7 @@ def encode(
         schema,
         chunk_length=chunk_length,
         chunk_width=chunk_width,
+        max_v_chunks=max_variant_chunks,
         worker_processes=worker_processes,
         show_progress=True,
     )
@@ -132,6 +162,9 @@ def convert_vcf(vcfs, out_path, chunk_length, chunk_width, verbose, worker_proce
 @click.argument("vcfs", nargs=-1, required=True)
 @click.argument("out_path", type=click.Path())
 def validate(vcfs, out_path):
+    """
+    Development only, do not use. Will be removed before release.
+    """
     # FIXME! Will silently not look at remaining VCFs
     vcf.validate(vcfs[0], out_path, show_progress=True)
 
@@ -158,7 +191,9 @@ def vcf2zarr():
 @verbose
 @chunk_length
 @chunk_width
-def convert_plink(in_path, out_path, verbose, worker_processes, chunk_length, chunk_width):
+def convert_plink(
+    in_path, out_path, verbose, worker_processes, chunk_length, chunk_width
+):
     """
     In development; DO NOT USE!
     """

bio2zarr/core.py

@@ -23,13 +23,15 @@
 )
 
 
-def chunk_aligned_slices(z, n):
+def chunk_aligned_slices(z, n, max_chunks=None):
     """
     Returns at n slices in the specified zarr array, aligned
     with its chunks
     """
     chunk_size = z.chunks[0]
     num_chunks = int(np.ceil(z.shape[0] / chunk_size))
+    if max_chunks is not None:
+        num_chunks = min(num_chunks, max_chunks)
     slices = []
     splits = np.array_split(np.arange(num_chunks), min(n, num_chunks))
     for split in splits:
@@ -132,7 +134,7 @@ class ProgressConfig:
     units: str = ""
     title: str = ""
     show: bool = False
-    poll_interval: float = 0.001
+    poll_interval: float = 0.01
 
 
 # NOTE: this approach means that we cannot have more than one
@@ -186,6 +188,7 @@ def __init__(self, worker_processes=1, progress_config=None):
         self.progress_thread = threading.Thread(
             target=self._update_progress_worker,
             name="progress-update",
+            daemon=True,  # Avoids deadlock on exit in awkward error conditions
         )
         self.progress_thread.start()
 

bio2zarr/vcf.py

@@ -1091,7 +1091,7 @@ def generate(pcvcf, chunk_length=None, chunk_width=None):
             chunk_width = 1000
         if chunk_length is None:
             chunk_length = 10_000
-
+        logger.info(f"Generating schema with chunks={chunk_length, chunk_width}")
         compressor = core.default_compressor.get_config()
 
         def fixed_field_spec(
@@ -1403,7 +1403,13 @@ def encode_filter_id(self, pcvcf, filter_names):
 
     @staticmethod
     def encode(
-        pcvcf, path, conversion_spec, *, worker_processes=1, show_progress=False
+        pcvcf,
+        path,
+        conversion_spec,
+        *,
+        worker_processes=1,
+        max_v_chunks=None,
+        show_progress=False,
     ):
         path = pathlib.Path(path)
         # TODO: we should do this as a future to avoid blocking
@@ -1425,7 +1431,15 @@ def encode(
 
         num_slices = max(1, worker_processes * 4)
         # Using POS arbitrarily to get the array slices
-        slices = core.chunk_aligned_slices(sgvcf.root["variant_position"], num_slices)
+        slices = core.chunk_aligned_slices(
+            sgvcf.root["variant_position"], num_slices, max_chunks=max_v_chunks
+        )
+        truncated = slices[-1][-1]
+        for array in sgvcf.root.values():
+            if array.attrs["_ARRAY_DIMENSIONS"][0] == "variants":
+                shape = list(array.shape)
+                shape[0] = truncated
+                array.resize(shape)
 
         chunked_1d = [
             col for col in conversion_spec.columns.values() if len(col.chunks) <= 1
@@ -1503,6 +1517,7 @@ def encode(
     schema_path=None,
     chunk_length=None,
     chunk_width=None,
+    max_v_chunks=None,
     worker_processes=1,
     show_progress=False,
 ):
@@ -1514,13 +1529,17 @@ def encode(
             chunk_width=chunk_width,
         )
     else:
-        # TODO checking that chunk_width and chunk_length are None
+        logger.info(f"Reading schema from {schema_path}")
+        if chunk_length is not None or chunk_width is not None:
+            raise ValueError("Cannot specify schema along with chunk sizes")
         with open(schema_path, "r") as f:
             schema = ZarrConversionSpec.fromjson(f.read())
+
     SgvcfZarr.encode(
         pcvcf,
         zarr_path,
         conversion_spec=schema,
+        max_v_chunks=max_v_chunks,
         worker_processes=worker_processes,
         show_progress=show_progress,
     )

tests/test_cli.py

@@ -66,6 +66,7 @@ def test_encode(self):
                 None,
                 chunk_length=None,
                 chunk_width=None,
+                max_v_chunks=None,
                 worker_processes=1,
                 show_progress=True,
             )

tests/test_core.py

@@ -72,6 +72,22 @@ def test_20_chunk_5(self, n, expected):
         result = core.chunk_aligned_slices(z, n)
         assert result == expected
 
+    @pytest.mark.parametrize(
+        ["n", "max_chunks", "expected"],
+        [
+            (1, 5, [(0, 20)]),
+            (1, 1, [(0, 5)]),
+            (2, 1, [(0, 5)]),
+            (3, 1, [(0, 5)]),
+            (2, 3, [(0, 10), (10, 15)]),
+            (2, 4, [(0, 10), (10, 20)]),
+        ],
+    )
+    def test_20_chunk_5_max_chunks(self, n, max_chunks, expected):
+        z = zarr.array(np.arange(20), chunks=5, dtype=int)
+        result = core.chunk_aligned_slices(z, n, max_chunks=max_chunks)
+        assert result == expected
+
     @pytest.mark.parametrize(
         ["n", "expected"],
         [

tests/test_vcf_examples.py

@@ -315,6 +315,16 @@ def test_full_pipeline(self, ds, tmp_path, worker_processes):
         ds2 = sg.load_dataset(out)
         xt.assert_equal(ds, ds2)
 
+    @pytest.mark.parametrize("max_v_chunks", [1, 2, 3])
+    @pytest.mark.parametrize("chunk_length", [1, 2, 3])
+    def test_max_v_chunks(self, ds, tmp_path, max_v_chunks, chunk_length):
+        exploded = tmp_path / "example.exploded"
+        vcf.explode([self.data_path], exploded)
+        out = tmp_path / "example.zarr"
+        vcf.encode(exploded, out, chunk_length=chunk_length, max_v_chunks=max_v_chunks)
+        ds2 = sg.load_dataset(out)
+        xt.assert_equal(ds.isel(variants=slice(None, chunk_length * max_v_chunks)), ds2)
+
     @pytest.mark.parametrize("worker_processes", [0, 1, 2])
     def test_worker_processes(self, ds, tmp_path, worker_processes):
         out = tmp_path / "example.vcf.zarr"