Add a --zarr-format CLI option

Author: tomwhite

bio2zarr/cli.py

@@ -43,6 +43,13 @@ def list_commands(self, ctx):
     "zarr_path", type=click.Path(exists=True, file_okay=False, dir_okay=True)
 )
 
+zarr_format = click.option(
+    "--zarr-format",
+    type=click.Choice([2, 3]),
+    default=None,
+    help="Zarr format version of output",
+)
+
 num_partitions = click.option(
     "-n",
     "--num-partitions",
@@ -500,6 +507,7 @@ def dencode_finalise(zarr_path, verbose, progress):
 @click.command(name="convert")
 @vcfs
 @new_zarr_path
+@zarr_format
 @force
 @variants_chunk_size
 @samples_chunk_size
@@ -510,6 +518,7 @@ def dencode_finalise(zarr_path, verbose, progress):
 def convert_vcf(
     vcfs,
     zarr_path,
+    zarr_format,
     force,
     variants_chunk_size,
     samples_chunk_size,
@@ -526,6 +535,7 @@ def convert_vcf(
     vcf_mod.convert(
         vcfs,
         zarr_path,
+        zarr_format=zarr_format,
         variants_chunk_size=variants_chunk_size,
         samples_chunk_size=samples_chunk_size,
         show_progress=progress,

bio2zarr/vcf.py

@@ -1654,6 +1654,7 @@ def convert(
     vcz_path=None,
     *,
     mode="r",
+    zarr_format=None,
     variants_chunk_size=None,
     samples_chunk_size=None,
     worker_processes=core.DEFAULT_WORKER_PROCESSES,
@@ -1719,6 +1720,7 @@ def convert(
             icf_path,
             vcz_path,
             mode=mode,
+            zarr_format=zarr_format,
             variants_chunk_size=variants_chunk_size,
             samples_chunk_size=samples_chunk_size,
             worker_processes=worker_processes,
@@ -1747,6 +1749,7 @@ def encode(
     worker_processes=core.DEFAULT_WORKER_PROCESSES,
     show_progress=False,
     mode="r",
+    zarr_format=None,
 ):
     with vcz.open_zarr(zarr_path, mode=mode) as zr:
         # Rough heuristic to split work up enough to keep utilisation high
@@ -1761,6 +1764,7 @@ def encode(
             local_alleles=local_alleles,
             ploidy=ploidy,
             max_variant_chunks=max_variant_chunks,
+            zarr_format=zarr_format,
         )
         vzw = vcz.VcfZarrWriter(IntermediateColumnarFormat, zr.dir)
         vzw.encode_all_partitions(
@@ -1787,6 +1791,7 @@ def encode_init(
     max_memory=None,
     worker_processes=core.DEFAULT_WORKER_PROCESSES,
     show_progress=False,
+    zarr_format=None,
 ):
     icf_store = IntermediateColumnarFormat(icf_path)
     if schema_path is None:
@@ -1805,7 +1810,7 @@ def encode_init(
         with open(schema_path) as f:
             schema_instance = vcz.VcfZarrSchema.fromjson(f.read())
     zarr_path = pathlib.Path(zarr_path)
-    vzw = vcz.VcfZarrWriter("icf", zarr_path)
+    vzw = vcz.VcfZarrWriter("icf", zarr_path, zarr_format=zarr_format)
     return vzw.init(
         icf_store,
         target_num_partitions=target_num_partitions,

bio2zarr/vcz.py

@@ -618,9 +618,10 @@ def encode(
 
 
 class VcfZarrWriter:
-    def __init__(self, source_type, path):
+    def __init__(self, source_type, path, zarr_format=None):
         self.source_type = source_type
         self.path = pathlib.Path(path)
+        self.zarr_format = zarr_format or zarr_utils.ZARR_FORMAT
         self.wip_path = self.path / "wip"
         self.arrays_path = self.wip_path / "arrays"
         self.partitions_path = self.wip_path / "partitions"
@@ -679,7 +680,10 @@ def init(
         )
 
         self.path.mkdir()
-        root = zarr.open(store=self.path, mode="a", **zarr_utils.ZARR_FORMAT_KWARGS)
+        zarr_format_kwargs = dict(
+            zarr_format=self.zarr_format or zarr_utils.ZARR_FORMAT
+        )
+        root = zarr.open(store=self.path, mode="a", **zarr_format_kwargs)
         root.attrs.update(
             {
                 "vcf_zarr_version": "0.4",
@@ -698,9 +702,7 @@ def init(
         self.wip_path.mkdir()
         self.arrays_path.mkdir()
         self.partitions_path.mkdir()
-        root = zarr.open(
-            store=self.arrays_path, mode="a", **zarr_utils.ZARR_FORMAT_KWARGS
-        )
+        root = zarr.open(store=self.arrays_path, mode="a", **zarr_format_kwargs)
 
         total_chunks = 0
         for field in self.schema.fields:
@@ -788,7 +790,7 @@ def init_array(self, root, schema, array_spec, variants_dim_size):
             else schema.defaults["compressor"]
         )
 
-        kwargs = dict(zarr_utils.ZARR_FORMAT_KWARGS)
+        kwargs = {}
         # see https://github.com/zarr-developers/zarr-python/issues/3197
         kwargs["fill_value"] = None
 
@@ -1050,7 +1052,9 @@ def finalise_array(self, name):
             if not src.exists():
                 # Needs test
                 raise ValueError(f"Partition {partition} of {name} does not exist")
-            zarr_utils.move_chunks(src, self.arrays_path, partition, name)
+            zarr_utils.move_chunks(
+                src, self.arrays_path, partition, name, self.zarr_format
+            )
         # Finally, once all the chunks have moved into the arrays dir,
         # we move it out of wip
         os.rename(self.arrays_path / name, self.path / name)

bio2zarr/zarr_utils.py

@@ -25,7 +25,6 @@
     ZARR_FORMAT = 2
 
 
-ZARR_FORMAT_KWARGS = dict(zarr_format=ZARR_FORMAT)
 # In zarr-python v3 strings are stored as string arrays (T) with itemsize 16
 STRING_DTYPE_NAME = "T"
 STRING_ITEMSIZE = 16
@@ -63,7 +62,7 @@
 }
 
 
-def make_compressor(config):
+def make_compressor(config, zarr_format):
     """Build a format-correct compressor from a numcodecs-style config dict.
 
     Returns a numcodecs codec under ZARR_FORMAT==2 and a zarr v3 codec under
@@ -74,7 +73,7 @@ def make_compressor(config):
         raise NotImplementedError(
             f"Only blosc compressors are supported, got {config!r}"
         )
-    if ZARR_FORMAT == 2:
+    if zarr_format == 2:
         return numcodecs.get_codec(config)
     return BloscCodec(
         cname=config["cname"],
@@ -123,11 +122,13 @@ def create_group_array(
     """Create an array within a group."""
     new_kwargs = {**kwargs}
     if compressor is not None:
-        new_kwargs["compressors"] = [make_compressor(compressor)]
+        new_kwargs["compressors"] = [
+            make_compressor(compressor, group.metadata.zarr_format)
+        ]
 
     # Zarr format v2 rejects dimension_names on create_array; we instead
     # write the xarray _ARRAY_DIMENSIONS attribute after the fact.
-    if ZARR_FORMAT == 3:
+    if group.metadata.zarr_format == 3:
         new_kwargs.pop("zarr_format", None)
         new_kwargs["dimension_names"] = dimension_names
 
@@ -139,7 +140,7 @@ def create_group_array(
         )
     else:
         array = group.create_array(name, shape=shape, dtype=dtype, **new_kwargs)
-    if ZARR_FORMAT == 2 and dimension_names is not None:
+    if group.metadata.zarr_format == 2 and dimension_names is not None:
         array.attrs["_ARRAY_DIMENSIONS"] = dimension_names
     return array
 
@@ -160,8 +161,10 @@ def create_empty_group_array(
     new_kwargs = {**kwargs}
     new_kwargs.pop("zarr_format", None)
     if compressor is not None:
-        new_kwargs["compressors"] = [make_compressor(compressor)]
-    if ZARR_FORMAT == 2:
+        new_kwargs["compressors"] = [
+            make_compressor(compressor, group.metadata.zarr_format)
+        ]
+    if group.metadata.zarr_format == 2:
         # Filter list is interpreted as numcodecs-style config dicts and
         # converted here so callers don't need to import numcodecs.
         # Zarr v2 also requires a VLenUTF8 filter to accompany "T"-dtype
@@ -184,7 +187,7 @@ def create_empty_group_array(
     array = group.create_array(
         name=name, shape=shape, dtype=dtype, chunks=chunks, **new_kwargs
     )
-    if ZARR_FORMAT == 2 and dimension_names is not None:
+    if group.metadata.zarr_format == 2 and dimension_names is not None:
         array.attrs["_ARRAY_DIMENSIONS"] = dimension_names
     return array
 
@@ -210,10 +213,10 @@ def get_compressor_config(array):
     raise TypeError(f"Unsupported compressor type: {type(compressor).__name__}")
 
 
-def move_chunks(src_path, dest_path, partition, name):
+def move_chunks(src_path, dest_path, partition, name, zarr_format):
     # Zarr v2 stores chunk files directly in the array directory; v3 places
     # them under a c/ subdirectory.
-    if ZARR_FORMAT == 2:
+    if zarr_format == 2:
         dest = dest_path / name
         chunk_files = [
             path for path in src_path.iterdir() if not path.name.startswith(".")

tests/test_cli.py

@@ -66,6 +66,7 @@
     show_progress=True,
     worker_processes=core.DEFAULT_WORKER_PROCESSES,
     local_alleles=False,
+    zarr_format=None,
 )
 
 DEFAULT_TSKIT_CONVERT_ARGS = dict(

tests/test_zarr_utils.py

@@ -133,15 +133,17 @@ def test_default_mode_read_only(self, tmp_path):
 
 
 class TestMakeCompressor:
-    def test_v2_returns_numcodecs(self, monkeypatch):
-        monkeypatch.setattr(zarr_utils, "ZARR_FORMAT", 2)
-        c = zarr_utils.make_compressor(zarr_utils.DEFAULT_COMPRESSOR_CONFIG)
+    def test_v2_returns_numcodecs(self):
+        c = zarr_utils.make_compressor(
+            zarr_utils.DEFAULT_COMPRESSOR_CONFIG, zarr_format=2
+        )
         assert isinstance(c, numcodecs.Blosc)
         assert c.get_config() == zarr_utils.DEFAULT_COMPRESSOR_CONFIG
 
-    def test_v3_returns_blosc_codec(self, monkeypatch):
-        monkeypatch.setattr(zarr_utils, "ZARR_FORMAT", 3)
-        c = zarr_utils.make_compressor(zarr_utils.DEFAULT_COMPRESSOR_CONFIG)
+    def test_v3_returns_blosc_codec(self):
+        c = zarr_utils.make_compressor(
+            zarr_utils.DEFAULT_COMPRESSOR_CONFIG, zarr_format=3
+        )
         assert isinstance(c, BloscCodec)
         assert c.cname.value == "zstd"
         assert c.clevel == 7
@@ -156,31 +158,36 @@ def test_v3_returns_blosc_codec(self, monkeypatch):
             (numcodecs.Blosc.BITSHUFFLE, BloscShuffle.bitshuffle),
         ],
     )
-    def test_v3_shuffle_mapping(self, monkeypatch, numcodecs_shuffle, v3_shuffle):
-        monkeypatch.setattr(zarr_utils, "ZARR_FORMAT", 3)
+    def test_v3_shuffle_mapping(self, numcodecs_shuffle, v3_shuffle):
         c = zarr_utils.make_compressor(
             {
                 "id": "blosc",
                 "cname": "zstd",
                 "clevel": 5,
                 "shuffle": numcodecs_shuffle,
-            }
+            },
+            zarr_format=3,
         )
         assert c.shuffle == v3_shuffle
         assert c.blocksize == 0
 
-    def test_v3_default_shuffle_when_omitted(self, monkeypatch):
-        monkeypatch.setattr(zarr_utils, "ZARR_FORMAT", 3)
-        c = zarr_utils.make_compressor({"id": "blosc", "cname": "lz4", "clevel": 1})
+    def test_v3_default_shuffle_when_omitted(self):
+        c = zarr_utils.make_compressor(
+            {"id": "blosc", "cname": "lz4", "clevel": 1}, zarr_format=3
+        )
         assert c.shuffle == BloscShuffle.shuffle
 
+    @pytest.mark.parametrize("zarr_format", [2, 3])
     def test_non_blosc_raises(self, zarr_format):
         with pytest.raises(NotImplementedError, match="Only blosc"):
-            zarr_utils.make_compressor({"id": "zlib", "level": 1})
+            zarr_utils.make_compressor(
+                {"id": "zlib", "level": 1}, zarr_format=zarr_format
+            )
 
+    @pytest.mark.parametrize("zarr_format", [2, 3])
     def test_missing_id_raises(self, zarr_format):
         with pytest.raises(NotImplementedError, match="Only blosc"):
-            zarr_utils.make_compressor({"cname": "zstd"})
+            zarr_utils.make_compressor({"cname": "zstd"}, zarr_format=zarr_format)
 
 
 class TestDefaultCompressorConstants:
@@ -576,8 +583,7 @@ class Stub:
 
 
 class TestMoveChunks:
-    def test_v2_layout(self, tmp_path, monkeypatch):
-        monkeypatch.setattr(zarr_utils, "ZARR_FORMAT", 2)
+    def test_v2_layout(self, tmp_path):
         src = tmp_path / "src"
         src.mkdir()
         (src / ".zarray").write_text("{}")
@@ -587,15 +593,14 @@ def test_v2_layout(self, tmp_path, monkeypatch):
         dest_root = tmp_path / "dest"
         (dest_root / "arr").mkdir(parents=True)
 
-        zarr_utils.move_chunks(src, dest_root, partition=0, name="arr")
+        zarr_utils.move_chunks(src, dest_root, partition=0, name="arr", zarr_format=2)
 
         assert (dest_root / "arr" / "0").read_text() == "chunk0"
         assert (dest_root / "arr" / "1").read_text() == "chunk1"
         # Hidden file is left behind.
         assert (src / ".zarray").exists()
 
-    def test_v3_layout(self, tmp_path, monkeypatch):
-        monkeypatch.setattr(zarr_utils, "ZARR_FORMAT", 3)
+    def test_v3_layout(self, tmp_path):
         src = tmp_path / "src"
         (src / "c").mkdir(parents=True)
         (src / "c" / "0").write_text("chunk0")
@@ -605,20 +610,19 @@ def test_v3_layout(self, tmp_path, monkeypatch):
         dest_root = tmp_path / "dest"
         (dest_root / "arr").mkdir(parents=True)
 
-        zarr_utils.move_chunks(src, dest_root, partition=0, name="arr")
+        zarr_utils.move_chunks(src, dest_root, partition=0, name="arr", zarr_format=3)
 
         assert (dest_root / "arr" / "c" / "0").read_text() == "chunk0"
         assert (dest_root / "arr" / "c" / "1").read_text() == "chunk1"
         assert (src / "c" / ".hidden").exists()
 
-    def test_v3_missing_c_directory(self, tmp_path, monkeypatch):
-        monkeypatch.setattr(zarr_utils, "ZARR_FORMAT", 3)
+    def test_v3_missing_c_directory(self, tmp_path):
         src = tmp_path / "src"
         src.mkdir()
 
         dest_root = tmp_path / "dest"
         (dest_root / "arr").mkdir(parents=True)
 
         # Should not raise even though src/c/ does not exist.
-        zarr_utils.move_chunks(src, dest_root, partition=0, name="arr")
+        zarr_utils.move_chunks(src, dest_root, partition=0, name="arr", zarr_format=3)
         assert list((dest_root / "arr" / "c").iterdir()) == []