Add iter_contig

Author: benjeffery

bio2zarr/vcf2zarr/icf.py

@@ -943,6 +943,19 @@ def iter_filters(self, start, stop):
                         ) from None
             yield filters
 
+    def iter_contig(self, start, stop):
+        source_field = self.fields["CHROM"]
+        lookup = {
+            contig.id: index for index, contig in enumerate(self.metadata.contigs)
+        }
+
+        for value in source_field.iter_values(start, stop):
+            # Note: because we are using the indexes to define the lookups
+            # and we always have an index, it seems that we the contig lookup
+            # will always succeed. However, if anyone ever does hit a KeyError
+            # here, please do open an issue with a reproducible example!
+            yield lookup[value[0]]
+
     def iter_field(self, field_name, shape, start, stop):
         source_field = self.fields[field_name]
         sanitiser = source_field.sanitiser_factory(shape)

bio2zarr/writer.py

@@ -557,18 +557,12 @@ def encode_filters_partition(self, partition_index):
         self.finalise_partition_array(partition_index, var_filter)
 
     def encode_contig_partition(self, partition_index):
-        lookup = {contig.id: index for index, contig in enumerate(self.schema.contigs)}
         contig = self.init_partition_array(partition_index, "variant_contig")
         partition = self.metadata.partitions[partition_index]
-        field = self.source.fields["CHROM"]
 
-        for value in field.iter_values(partition.start, partition.stop):
+        for contig_index in self.source.iter_contig(partition.start, partition.stop):
             j = contig.next_buffer_row()
-            # Note: because we are using the indexes to define the lookups
-            # and we always have an index, it seems that we the contig lookup
-            # will always succeed. However, if anyone ever does hit a KeyError
-            # here, please do open an issue with a reproducible example!
-            contig.buff[j] = lookup[value[0]]
+            contig.buff[j] = contig_index
 
         self.finalise_partition_array(partition_index, contig)