Switch partition_into_regions to return an iterator

Author: jeromekelleher

bio2zarr/cli.py

@@ -518,10 +518,10 @@ def plink2zarr():
 @click.command
 @version
 @click.argument("vcf_path", type=click.Path())
-@click.option("-i", "--index", type=click.Path(), default=None)
 @click.option("-n", "--num-parts", type=int, default=None)
 # @click.option("-s", "--part-size", type=int, default=None)
-def vcfpartition(vcf_path, index, num_parts):
-    indexed_vcf = vcf_utils.IndexedVcf(vcf_path, index)
+def vcfpartition(vcf_path, num_parts):
+    indexed_vcf = vcf_utils.IndexedVcf(vcf_path)
     regions = indexed_vcf.partition_into_regions(num_parts=num_parts)
-    click.echo("\n".join(map(str, regions)))
+    for region in regions:
+        click.echo(region)

bio2zarr/vcf2zarr/icf.py

@@ -263,9 +263,6 @@ def scan_vcf(path, target_num_partitions):
         )
 
         regions = indexed_vcf.partition_into_regions(num_parts=target_num_partitions)
-        logger.info(
-            f"Split {path} into {len(regions)} regions (target={target_num_partitions})"
-        )
         for region in regions:
             metadata.partitions.append(
                 VcfPartition(
@@ -275,6 +272,10 @@ def scan_vcf(path, target_num_partitions):
                     region=region,
                 )
             )
+        logger.info(
+            f"Split {path} into {len(metadata.partitions)} "
+            f"partitions target={target_num_partitions})"
+        )
         core.update_progress(1)
         return metadata, vcf.raw_header
 

bio2zarr/vcf_utils.py

@@ -450,15 +450,16 @@ def variants(self, region):
     def _filter_empty_and_refine(self, regions):
         """
         Return all regions in the specified list that have one or more records,
-        and refine the start coordinate of the region to be the actual first coord
+        and refine the start coordinate of the region to be the actual first coord.
+
+        Because this is a relatively expensive operation requiring seeking around
+        the file, we return the results as an iterator.
         """
-        ret = []
         for region in regions:
             var = next(self.variants(region), None)
             if var is not None:
                 region.start = var.POS
-                ret.append(region)
-        return ret
+                yield region
 
     def partition_into_regions(
         self,

tests/test_vcf_utils.py

@@ -95,7 +95,7 @@ def test_contig_record_counts(self, index_file, expected):
     )
     def test_partition_into_one_part(self, index_file, expected):
         indexed_vcf = self.get_instance(index_file)
-        regions = indexed_vcf.partition_into_regions(num_parts=1)
+        regions = list(indexed_vcf.partition_into_regions(num_parts=1))
         assert all(isinstance(r, vcf_utils.Region) for r in regions)
         assert [str(r) for r in regions] == expected
 
@@ -120,7 +120,7 @@ def test_partition_into_one_part(self, index_file, expected):
     )
     def test_partition_into_max_parts(self, index_file, num_expected, total_records):
         indexed_vcf = self.get_instance(index_file)
-        regions = indexed_vcf.partition_into_regions(num_parts=1000)
+        regions = list(indexed_vcf.partition_into_regions(num_parts=1000))
         assert all(isinstance(r, vcf_utils.Region) for r in regions)
         # print(regions)
         assert len(regions) == num_expected
@@ -151,7 +151,7 @@ def test_partition_into_max_parts(self, index_file, num_expected, total_records)
     @pytest.mark.parametrize("num_parts", [2, 3, 4, 5, 16, 33])
     def test_partition_into_n_parts(self, index_file, total_records, num_parts):
         indexed_vcf = self.get_instance(index_file)
-        regions = indexed_vcf.partition_into_regions(num_parts=num_parts)
+        regions = list(indexed_vcf.partition_into_regions(num_parts=num_parts))
         assert all(isinstance(r, vcf_utils.Region) for r in regions)
         part_variant_counts = np.array(
             [indexed_vcf.count_variants(region) for region in regions]
@@ -161,7 +161,7 @@ def test_partition_into_n_parts(self, index_file, total_records, num_parts):
 
     def test_tabix_multi_chrom_bug(self):
         indexed_vcf = self.get_instance("multi_contig.vcf.gz.tbi")
-        regions = indexed_vcf.partition_into_regions(num_parts=10)
+        regions = list(indexed_vcf.partition_into_regions(num_parts=10))
         # An earlier version of the code returned this, i.e. with a duplicate
         # for 4 with end coord of 0
         # ["0:1-", "1", "2", "3", "4:1-0", "4:1-"]
@@ -185,7 +185,9 @@ def test_tabix_multi_chrom_bug(self):
     )
     def test_target_part_size(self, target_part_size, filename):
         indexed_vcf = self.get_instance(filename)
-        regions = indexed_vcf.partition_into_regions(target_part_size=target_part_size)
+        regions = list(
+            indexed_vcf.partition_into_regions(target_part_size=target_part_size)
+        )
         assert len(regions) == 5
         part_variant_counts = [indexed_vcf.count_variants(region) for region in regions]
         assert part_variant_counts == [3450, 3869, 4525, 7041, 1025]
@@ -197,7 +199,7 @@ def test_partition_invalid_arguments(self):
         with pytest.raises(
             ValueError, match=r"One of num_parts or target_part_size must be specified"
         ):
-            indexed_vcf.partition_into_regions()
+            list(indexed_vcf.partition_into_regions())
 
         with pytest.raises(
             ValueError,