Add first_dim_iter utility to workaround slow array iteration in zarr-python v3

Author: tomwhite

bio2zarr/vcf2zarr/verification.py

@@ -4,6 +4,8 @@
 import tqdm
 import zarr
 
+from bio2zarr.zarr_utils import first_dim_iter
+
 from .. import constants
 
 
@@ -152,7 +154,7 @@ def verify(vcf_path, zarr_path, show_progress=False):
     vid = root["variant_id"][:]
     call_genotype = None
     if "call_genotype" in root and root["call_genotype"].size > 0:
-        call_genotype = iter(root["call_genotype"])
+        call_genotype = first_dim_iter(root["call_genotype"])
 
     vcf = cyvcf2.VCF(vcf_path)
     format_headers = {}
@@ -170,12 +172,16 @@ def verify(vcf_path, zarr_path, show_progress=False):
             vcf_name = colname.split("_", 1)[1]
             vcf_type = format_headers[vcf_name]["Type"]
             vcf_number = format_headers[vcf_name]["Number"]
-            format_fields[vcf_name] = vcf_type, vcf_number, iter(root[colname])
+            format_fields[vcf_name] = (
+                vcf_type,
+                vcf_number,
+                first_dim_iter(root[colname]),
+            )
         if colname.startswith("variant"):
             name = colname.split("_", 1)[1]
             if name.isupper():
                 vcf_type = info_headers[name]["Type"]
-                info_fields[name] = vcf_type, iter(root[colname])
+                info_fields[name] = vcf_type, first_dim_iter(root[colname])
 
     first_pos = next(vcf).POS
     start_index = np.searchsorted(pos, first_pos)

bio2zarr/zarr_utils.py

@@ -11,3 +11,9 @@ def zarr_v3() -> bool:
     ZARR_FORMAT_KWARGS = dict(zarr_format=2)
 else:
     ZARR_FORMAT_KWARGS = dict()
+
+
+# See discussion in https://github.com/zarr-developers/zarr-python/issues/2529
+def first_dim_iter(z):
+    for chunk in range(z.cdata_shape[0]):
+        yield from z.blocks[chunk]