Merge pull request #211 from jeromekelleher/final-vcfpartition-version

Final vcfpartition version

Author: jeromekelleher

bio2zarr/cli.py

@@ -518,10 +518,42 @@ 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)
-    regions = indexed_vcf.partition_into_regions(num_parts=num_parts)
-    click.echo("\n".join(map(str, regions)))
+@verbose
+@click.option(
+    "-n",
+    "--num-parts",
+    type=int,
+    default=None,
+    help="Target number of partitions to split the VCF into",
+)
+@click.option(
+    "-s",
+    "--part-size",
+    type=str,
+    default=None,
+    help="Target (compressed) size of VCF partitions, e.g. 100KB, 10MiB, 1G.",
+)
+def vcfpartition(vcf_path, verbose, num_parts, part_size):
+    """
+    Output bcftools region strings that partition an indexed VCF/BCF file
+    into either an approximate number of parts (-n), or parts of approximately
+    a given size (-s). One of -n or -s must be supplied.
+
+    Note that both the number of partitions and sizes are a target, and the
+    returned number of partitions may not exactly correspond. In particular,
+    there is a maximum level of granularity determined by the associated index
+    which cannot be exceeded.
+
+    Note also that the partitions returned may vary considerably in the number
+    of records that they contain.
+    """
+    setup_logging(verbose)
+    if num_parts is None and part_size is None:
+        raise click.UsageError("Either --num-parts or --part-size must be specified")
+
+    indexed_vcf = vcf_utils.IndexedVcf(vcf_path)
+    regions = indexed_vcf.partition_into_regions(
+        num_parts=num_parts, target_part_size=part_size
+    )
+    for region in regions:
+        click.echo(f"{region}\t{vcf_path}")

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

@@ -1,5 +1,6 @@
 import contextlib
 import gzip
+import logging
 import os
 import pathlib
 import struct
@@ -13,6 +14,8 @@
 
 from bio2zarr.typing import PathType
 
+logger = logging.getLogger(__name__)
+
 CSI_EXTENSION = ".csi"
 TABIX_EXTENSION = ".tbi"
 TABIX_LINEAR_INDEX_INTERVAL_SIZE = 1 << 14  # 16kb interval size
@@ -411,6 +414,7 @@ def __init__(self, vcf_path, index_path=None):
             raise ValueError("Only .tbi or .csi indexes are supported.")
         self.vcf = cyvcf2.VCF(vcf_path)
         self.vcf.set_index(str(self.index_path))
+        logger.debug(f"Loaded {vcf_path} with index {self.index_path}")
         self.sequence_names = None
         if self.index_type == "csi":
             # Determine the file-type based on the "aux" field.
@@ -450,15 +454,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,
@@ -490,7 +495,7 @@ def partition_into_regions(
             target_part_size_bytes = file_length // num_parts
         elif target_part_size_bytes is not None:
             num_parts = ceildiv(file_length, target_part_size_bytes)
-        part_lengths = np.array([i * target_part_size_bytes for i in range(num_parts)])
+        part_lengths = target_part_size_bytes * np.arange(num_parts, dtype=int)
         file_offsets, region_contig_indexes, region_positions = self.index.offsets()
 
         # Search the file offsets to find which indexes the part lengths fall at

tests/test_cli.py

@@ -697,21 +697,55 @@ def test_convert(self, tmp_path):
 
 
 class TestVcfPartition:
+    path = "tests/data/vcf/NA12878.prod.chr20snippet.g.vcf.gz"
+
     def test_num_parts(self):
-        path = "tests/data/vcf/NA12878.prod.chr20snippet.g.vcf.gz"
+        runner = ct.CliRunner(mix_stderr=False)
+        result = runner.invoke(
+            cli.vcfpartition, [self.path, "-n", "5"], catch_exceptions=False
+        )
+        assert result.stderr == ""
+        assert result.exit_code == 0
+        assert list(result.stdout.splitlines()) == [
+            "20:60001-278528\ttests/data/vcf/NA12878.prod.chr20snippet.g.vcf.gz",
+            "20:278529-442368\ttests/data/vcf/NA12878.prod.chr20snippet.g.vcf.gz",
+            "20:442381-638976\ttests/data/vcf/NA12878.prod.chr20snippet.g.vcf.gz",
+            "20:638982-819200\ttests/data/vcf/NA12878.prod.chr20snippet.g.vcf.gz",
+            "20:819201-\ttests/data/vcf/NA12878.prod.chr20snippet.g.vcf.gz",
+        ]
 
+    def test_part_size(self):
         runner = ct.CliRunner(mix_stderr=False)
         result = runner.invoke(
-            cli.vcfpartition, [path, "-n", "5"], catch_exceptions=False
+            cli.vcfpartition, [self.path, "-s", "512K"], catch_exceptions=False
         )
+        assert result.stderr == ""
+        assert result.exit_code == 0
         assert list(result.stdout.splitlines()) == [
-            "20:60001-278528",
-            "20:278529-442368",
-            "20:442381-638976",
-            "20:638982-819200",
-            "20:819201-",
+            "20:60001-212992\ttests/data/vcf/NA12878.prod.chr20snippet.g.vcf.gz",
+            "20:213070-327680\ttests/data/vcf/NA12878.prod.chr20snippet.g.vcf.gz",
+            "20:327695-442368\ttests/data/vcf/NA12878.prod.chr20snippet.g.vcf.gz",
+            "20:442381-557056\ttests/data/vcf/NA12878.prod.chr20snippet.g.vcf.gz",
+            "20:557078-688128\ttests/data/vcf/NA12878.prod.chr20snippet.g.vcf.gz",
+            "20:688129-802816\ttests/data/vcf/NA12878.prod.chr20snippet.g.vcf.gz",
+            "20:802817-933888\ttests/data/vcf/NA12878.prod.chr20snippet.g.vcf.gz",
+            "20:933890-\ttests/data/vcf/NA12878.prod.chr20snippet.g.vcf.gz",
         ]
 
+    def test_no_part_spec(self):
+        runner = ct.CliRunner(mix_stderr=False)
+        result = runner.invoke(cli.vcfpartition, [self.path], catch_exceptions=False)
+        assert result.exit_code != 0
+        assert result.stdout == ""
+        assert len(result.stderr) > 0
+
+    def test_no_args(self):
+        runner = ct.CliRunner(mix_stderr=False)
+        result = runner.invoke(cli.vcfpartition, [], catch_exceptions=False)
+        assert result.exit_code != 0
+        assert result.stdout == ""
+        assert len(result.stderr) > 0
+
 
 @pytest.mark.parametrize(
     "cmd", [main.bio2zarr, cli.vcf2zarr_main, cli.plink2zarr, cli.vcfpartition]

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,