Remove duplicate indexer, fix zarr v3

Author: benjeffery

bio2zarr/plink.py

@@ -82,7 +82,7 @@ def generate_schema(
         schema.ZarrArraySpec.new(
             vcf_field=None,
             name="variant_allele",
-            dtype="str",
+            dtype="O",
             shape=[m, 2],
             dimensions=["variants", "alleles"],
             chunks=[variants_chunk_size, 2],
@@ -151,13 +151,10 @@ def convert(
         PlinkFormat(bed_path),
         target_num_partitions=target_num_partitions,
         schema=schema_instance,
-        #        dimension_separator=None,
-        #       max_variant_chunks=None
     )
     vzw.encode_all_partitions(
         worker_processes=worker_processes,
         show_progress=show_progress,
-        # max_memory=None,
     )
     vzw.finalise(show_progress)
 

bio2zarr/writer.py

@@ -641,60 +641,8 @@ def finalise(self, show_progress=False):
     def create_index(self):
         """Create an index to support efficient region queries."""
 
-        root = zarr.open_group(store=self.path, mode="r+")
-
-        contig = root["variant_contig"]
-        pos = root["variant_position"]
-        length = root["variant_length"]
-
-        assert contig.cdata_shape == pos.cdata_shape
-
-        index = []
-
-        logger.info("Creating region index")
-        for v_chunk in range(pos.cdata_shape[0]):
-            c = contig.blocks[v_chunk]
-            p = pos.blocks[v_chunk]
-            e = p + length.blocks[v_chunk] - 1
-
-            # create a row for each contig in the chunk
-            d = np.diff(c, append=-1)
-            c_start_idx = 0
-            for c_end_idx in np.nonzero(d)[0]:
-                assert c[c_start_idx] == c[c_end_idx]
-                index.append(
-                    (
-                        v_chunk,  # chunk index
-                        c[c_start_idx],  # contig ID
-                        p[c_start_idx],  # start
-                        p[c_end_idx],  # end
-                        np.max(e[c_start_idx : c_end_idx + 1]),  # max end
-                        c_end_idx - c_start_idx + 1,  # num records
-                    )
-                )
-                c_start_idx = c_end_idx + 1
-
-        index = np.array(index, dtype=pos.dtype)
-        kwargs = {}
-        if not zarr_utils.zarr_v3():
-            kwargs["dimension_separator"] = self.metadata.dimension_separator
-        array = root.array(
-            "region_index",
-            data=index,
-            shape=index.shape,
-            chunks=index.shape,
-            dtype=index.dtype,
-            compressor=numcodecs.Blosc("zstd", clevel=9, shuffle=0),
-            fill_value=None,
-            **kwargs,
-        )
-        array.attrs["_ARRAY_DIMENSIONS"] = [
-            "region_index_values",
-            "region_index_fields",
-        ]
-
-        logger.info("Consolidating Zarr metadata")
-        zarr.consolidate_metadata(self.path)
+        indexer = VcfZarrIndexer(self.path)
+        indexer.create_index()
 
     ######################
     # encode_all_partitions