Add --target-uncompressed-size to the rechunk CLI command (#100)

Author: tomwhite

tests/test_cli.py

@@ -170,7 +170,6 @@ def test_invalid_zarr_backend_storage_is_rejected(command, command_args):
     [
         ["append", "only-one-store"],
         ["normalise", "left", "right"],
-        ["rechunk", "store", "variant_contig"],
         ["remove", "store"],
         ["copy-store-to-icechunk", "only-one-store"],
     ],
@@ -228,15 +227,22 @@ def test_remove_cli_updates_vcz_store(tmp_path):
     check_removed_sample(vcz, "NA00002")
 
 
-def test_rechunk_cli_passes_arguments(monkeypatch):
-    seen = {}
+def fake_rechunk(
+    vcz,
+    variants_array_name,
+    variants_chunk_size=None,
+    *,
+    target_uncompressed_size_bytes=None,
+    backend_storage=None,
+):
+    fake_rechunk.seen = {
+        "args": (vcz, variants_array_name, variants_chunk_size),
+        "target_uncompressed_size_bytes": target_uncompressed_size_bytes,
+        "backend_storage": backend_storage,
+    }
 
-    def fake_rechunk(
-        vcz, variants_array_name, variants_chunk_size, *, backend_storage=None
-    ):
-        seen["args"] = (vcz, variants_array_name, variants_chunk_size)
-        seen["backend_storage"] = backend_storage
 
+def test_rechunk_cli_passes_arguments(monkeypatch):
     monkeypatch.setattr(cli, "rechunk_function", fake_rechunk)
 
     runner = ct.CliRunner()
@@ -247,12 +253,63 @@ def fake_rechunk(
     )
 
     assert result.exit_code == 0
-    assert seen == {
+    assert fake_rechunk.seen == {
         "args": ("store", "variant_contig", 4),
+        "target_uncompressed_size_bytes": None,
         "backend_storage": "icechunk",
     }
 
 
+def test_rechunk_cli_passes_target_uncompressed_size(monkeypatch):
+    monkeypatch.setattr(cli, "rechunk_function", fake_rechunk)
+
+    runner = ct.CliRunner()
+    result = runner.invoke(
+        cli.vczstore_main,
+        ["rechunk", "--target-uncompressed-size", "100MB", "store", "variant_contig"],
+        catch_exceptions=False,
+    )
+
+    assert result.exit_code == 0
+    assert fake_rechunk.seen == {
+        "args": ("store", "variant_contig", None),
+        "target_uncompressed_size_bytes": 100_000_000,
+        "backend_storage": None,
+    }
+
+
+def test_rechunk_cli_rejects_neither_size_option():
+    runner = ct.CliRunner()
+    result = runner.invoke(cli.vczstore_main, ["rechunk", "store", "variant_contig"])
+
+    assert result.exit_code == 2
+    assert (
+        "Must specify either VARIANTS_CHUNK_SIZE or --target-uncompressed-size"
+        in result.output
+    )
+
+
+def test_rechunk_cli_rejects_both_size_options():
+    runner = ct.CliRunner()
+    result = runner.invoke(
+        cli.vczstore_main,
+        [
+            "rechunk",
+            "--target-uncompressed-size",
+            "100MB",
+            "store",
+            "variant_contig",
+            "4",
+        ],
+    )
+
+    assert result.exit_code == 2
+    assert (
+        "Cannot specify both VARIANTS_CHUNK_SIZE and --target-uncompressed-size"
+        in result.output
+    )
+
+
 @pytest.mark.parametrize(
     ("backend_storage", "zarr_format"),
     [(None, None), ("icechunk", 3)],

tests/test_rechunk.py

@@ -1,8 +1,10 @@
+import numpy as np
 import pytest
+import zarr
 from vcztools.utils import open_zarr
 
 from tests.utils import make_vcz
-from vczstore.rechunk import rechunk
+from vczstore.rechunk import compute_variants_chunk_size_from_target, rechunk
 
 
 def make_simple_vcz(variants_chunk_size=2):
@@ -48,3 +50,44 @@ def test_rechunk_not_multiple_of_min_chunk_size():
     vcz = make_simple_vcz()
     with pytest.raises(ValueError, match="not an exact multiple"):
         rechunk(vcz, "variant_contig", 3)
+
+
+def test_rechunk_with_target_uncompressed_size():
+    vcz = make_simple_vcz()
+
+    rechunk(vcz, "variant_contig", target_uncompressed_size_bytes=16)
+
+    root = open_zarr(vcz)
+    # variant_contig is int32 (4 bytes), 1D; bytes_per_variant_chunk=4, n=4,
+    # rounded to multiple of 2 → 4
+    assert root["variant_contig"].chunks[0] == 4
+    assert root["call_genotype"].chunks[0] == 2
+
+
+@pytest.mark.parametrize(
+    ("dtype", "extra_chunks", "target_bytes", "min_chunk_size", "expected"),
+    [
+        # 1D int32: bytes_per_variant_chunk=4, target=16 → n=4, multiple of 2 → 4
+        (np.int32, (), 16, 2, 4),
+        # 1D int32: target=20 → n=5, rounded down to multiple of 2 → 4
+        (np.int32, (), 20, 2, 4),
+        # 1D int32: target too small → clamp to min_chunk_size
+        (np.int32, (), 1, 2, 2),
+        # 2D int8 with extra chunk dim 3: bytes_per_variant_chunk=3,
+        # target=12 → n=4, multiple of 2 → 4
+        (np.int8, (3,), 12, 2, 4),
+    ],
+)
+def test_compute_variants_chunk_size_from_target(
+    dtype, extra_chunks, target_bytes, min_chunk_size, expected
+):
+    store = zarr.storage.MemoryStore()
+    root = zarr.open_group(store=store)
+    shape = (10,) + tuple(c * 2 for c in extra_chunks)
+    chunks = (2,) + extra_chunks
+    arr = root.create_array(name="x", data=np.zeros(shape, dtype=dtype), chunks=chunks)
+
+    result = compute_variants_chunk_size_from_target(arr, target_bytes, min_chunk_size)
+
+    assert result == expected
+    assert result % min_chunk_size == 0

vczstore/cli.py

@@ -6,6 +6,7 @@
 from vczstore.normalise import normalise as normalise_function
 from vczstore.rechunk import rechunk as rechunk_function
 from vczstore.remove import remove as remove_function
+from vczstore.utils import parse_size
 
 
 class NaturalOrderGroup(click.Group):
@@ -155,18 +156,54 @@ def normalise(
 @click.command()
 @click.argument("vcz", type=click.Path())
 @click.argument("variants_array_name", type=str)
-@click.argument("variants_chunk_size", type=click.IntRange(min=1))
+@click.argument(
+    "variants_chunk_size", type=click.IntRange(min=1), required=False, default=None
+)
+@click.option(
+    "--target-uncompressed-size",
+    type=str,
+    default=None,
+    help=(
+        "Target uncompressed chunk size (e.g. '100MB'). "
+        "Alternative to VARIANTS_CHUNK_SIZE."
+    ),
+)
 @verbose
 @backend_storage
-def rechunk(vcz, variants_array_name, variants_chunk_size, verbose, backend_storage):
+def rechunk(
+    vcz,
+    variants_array_name,
+    variants_chunk_size,
+    target_uncompressed_size,
+    verbose,
+    backend_storage,
+):
     """Rechunk a variants array with a larger variants chunk size that is
     an exact multiple of the min variants chunk size"""
     setup_logging(verbose)
+    if variants_chunk_size is not None and target_uncompressed_size is not None:
+        raise click.UsageError(
+            "Cannot specify both VARIANTS_CHUNK_SIZE and --target-uncompressed-size"
+        )
+    if variants_chunk_size is None and target_uncompressed_size is None:
+        raise click.UsageError(
+            "Must specify either VARIANTS_CHUNK_SIZE or --target-uncompressed-size"
+        )
+    target_bytes = None
+    if target_uncompressed_size is not None:
+        try:
+            target_bytes = parse_size(target_uncompressed_size)
+        except ValueError as e:
+            raise click.BadParameter(
+                f"Cannot parse size: {target_uncompressed_size!r}",
+                param_hint="'--target-uncompressed-size'",
+            ) from e
     call_or_error(
         rechunk_function,
         vcz,
         variants_array_name,
         variants_chunk_size,
+        target_uncompressed_size_bytes=target_bytes,
         backend_storage=backend_storage,
     )
 

vczstore/rechunk.py

@@ -1,3 +1,5 @@
+import math
+
 from bio2zarr.zarr_utils import create_group_array, get_compressor_config
 from vcztools.utils import array_dims, open_zarr
 
@@ -11,8 +13,9 @@
 def rechunk(
     vcz,
     variants_array_name,
-    variants_chunk_size,
+    variants_chunk_size=None,
     *,
+    target_uncompressed_size_bytes=None,
     backend_storage=None,
 ):
     """Rechunk a variants array with a larger variants chunk size that is
@@ -30,6 +33,12 @@ def rechunk(
             )
 
         min_chunk_size = compute_min_variants_chunk_size(root)
+
+        if target_uncompressed_size_bytes is not None:
+            variants_chunk_size = compute_variants_chunk_size_from_target(
+                arr, target_uncompressed_size_bytes, min_chunk_size
+            )
+
         if variants_chunk_size % min_chunk_size != 0:
             raise ValueError(
                 f"variants_chunk_size={variants_chunk_size} is not an exact multiple "
@@ -55,3 +64,13 @@ def rechunk(
             compressor=get_compressor_config(arr),
             dimension_names=array_dims(arr),
         )
+
+
+def compute_variants_chunk_size_from_target(
+    arr, target_bytes: int, min_chunk_size: int
+) -> int:
+    """Compute the variants chunk size giving chunks of approximately
+    target_bytes uncompressed."""
+    bytes_per_variant_chunk = math.prod(arr.chunks[1:]) * arr.dtype.itemsize
+    n = target_bytes // bytes_per_variant_chunk
+    return max((n // min_chunk_size) * min_chunk_size, min_chunk_size)

vczstore/utils.py

@@ -17,6 +17,33 @@
 )
 
 
+def parse_size(s: str) -> int:
+    """Parse a human-readable size string (e.g. '100MB', '2GiB') to bytes."""
+    suffixes = [
+        ("TiB", 1024**4),
+        ("GiB", 1024**3),
+        ("MiB", 1024**2),
+        ("KiB", 1024),
+        ("TB", 1000**4),
+        ("GB", 1000**3),
+        ("MB", 1000**2),
+        ("KB", 1000),
+        ("B", 1),
+    ]
+    upper = s.strip().upper()
+    for suffix, factor in suffixes:
+        if upper.endswith(suffix):
+            num = s[: len(s) - len(suffix)].strip()
+            try:
+                return int(float(num) * factor)
+            except ValueError as e:
+                raise ValueError(f"Cannot parse size: {s!r}") from e
+    try:
+        return int(s.strip())
+    except ValueError as e:
+        raise ValueError(f"Cannot parse size: {s!r}") from e
+
+
 def missing_val(arr):
     if arr.dtype.kind == "i":
         return INT_MISSING