Changes to run on Zarr v3 (but always using v2 for Zarr format)

Author: tomwhite

bio2zarr/vcf2zarr/vcz.py

@@ -12,6 +12,7 @@
 import numcodecs
 import numpy as np
 import zarr
+from packaging.version import Version
 
 from .. import constants, core, provenance
 from . import icf
@@ -470,6 +471,17 @@ class VcfZarrWriteSummary(core.JsonDataclass):
     max_encoding_memory: str
 
 
+def _zarr_v3() -> bool:
+    return Version(zarr.__version__).major >= 3
+
+
+if _zarr_v3():
+    # Use zarr format v2 even when running with zarr-python v3
+    ZARR_FORMAT_KWARGS = dict(zarr_format=2)
+else:
+    ZARR_FORMAT_KWARGS = dict()
+
+
 class VcfZarrWriter:
     def __init__(self, path):
         self.path = pathlib.Path(path)
@@ -532,7 +544,7 @@ def init(
         )
 
         self.path.mkdir()
-        root = zarr.open(store=self.path, mode="a")
+        root = zarr.open(store=self.path, mode="a", **ZARR_FORMAT_KWARGS)
         root.attrs.update(
             {
                 "vcf_zarr_version": "0.2",
@@ -548,7 +560,7 @@ def init(
         self.wip_path.mkdir()
         self.arrays_path.mkdir()
         self.partitions_path.mkdir()
-        root = zarr.open(store=self.arrays_path, mode="a")
+        root = zarr.open(store=self.arrays_path, mode="a", **ZARR_FORMAT_KWARGS)
 
         total_chunks = 0
         for field in self.schema.fields:
@@ -572,7 +584,8 @@ def encode_samples(self, root):
             raise ValueError("Subsetting or reordering samples not supported currently")
         array = root.array(
             "sample_id",
-            [sample.id for sample in self.schema.samples],
+            data=[sample.id for sample in self.schema.samples],
+            shape=len(self.schema.samples),
             dtype="str",
             compressor=DEFAULT_ZARR_COMPRESSOR,
             chunks=(self.schema.samples_chunk_size,),
@@ -583,15 +596,17 @@ def encode_samples(self, root):
     def encode_contig_id(self, root):
         array = root.array(
             "contig_id",
-            [contig.id for contig in self.schema.contigs],
+            data=[contig.id for contig in self.schema.contigs],
+            shape=len(self.schema.contigs),
             dtype="str",
             compressor=DEFAULT_ZARR_COMPRESSOR,
         )
         array.attrs["_ARRAY_DIMENSIONS"] = ["contigs"]
         if all(contig.length is not None for contig in self.schema.contigs):
             array = root.array(
                 "contig_length",
-                [contig.length for contig in self.schema.contigs],
+                data=[contig.length for contig in self.schema.contigs],
+                shape=len(self.schema.contigs),
                 dtype=np.int64,
                 compressor=DEFAULT_ZARR_COMPRESSOR,
             )
@@ -602,7 +617,8 @@ def encode_filter_id(self, root):
         # https://github.com/sgkit-dev/vcf-zarr-spec/issues/19
         array = root.array(
             "filter_id",
-            [filt.id for filt in self.schema.filters],
+            data=[filt.id for filt in self.schema.filters],
+            shape=len(self.schema.filters),
             dtype="str",
             compressor=DEFAULT_ZARR_COMPRESSOR,
         )
@@ -616,14 +632,15 @@ def init_array(self, root, array_spec, variants_dim_size):
         # Truncate the variants dimension is max_variant_chunks was specified
         shape[0] = variants_dim_size
         a = root.empty(
-            array_spec.name,
+            name=array_spec.name,
             shape=shape,
             chunks=array_spec.chunks,
             dtype=array_spec.dtype,
             compressor=numcodecs.get_codec(array_spec.compressor),
             filters=[numcodecs.get_codec(filt) for filt in array_spec.filters],
             object_codec=object_codec,
             dimension_separator=self.metadata.dimension_separator,
+            **ZARR_FORMAT_KWARGS,
         )
         a.attrs.update(
             {
@@ -688,7 +705,11 @@ def init_partition_array(self, partition_index, name):
         # Overwrite any existing WIP files
         wip_path = self.wip_partition_array_path(partition_index, name)
         shutil.copytree(src, wip_path, dirs_exist_ok=True)
-        wip_root = zarr.open(store=self.wip_partition_path(partition_index), mode="a")
+        wip_root = zarr.open(
+            store=self.wip_partition_path(partition_index),
+            mode="a",
+            **ZARR_FORMAT_KWARGS,
+        )
         array = wip_root[name]
         logger.debug(f"Opened empty array {array.name} <{array.dtype}> @ {wip_path}")
         return array