Add vcf-partition command and isolate Tabix issue

Author: jeromekelleher

bio2zarr/__main__.py

@@ -14,6 +14,7 @@ def bio2zarr():
 # up in the right way.
 bio2zarr.add_command(cli.vcf2zarr)
 bio2zarr.add_command(cli.plink2zarr)
+bio2zarr.add_command(cli.vcf_partition)
 
 if __name__ == "__main__":
     bio2zarr()

bio2zarr/cli.py

@@ -3,6 +3,7 @@
 import coloredlogs
 
 from . import vcf
+from . import vcf_utils
 from . import plink
 from . import provenance
 
@@ -15,6 +16,7 @@
 
 version = click.version_option(version=provenance.__version__)
 
+
 # Note: logging hasn't been implemented in the code at all, this is just
 # a first pass to try out some ways of doing things to see what works.
 def setup_logging(verbosity):
@@ -155,3 +157,15 @@ def plink2zarr():
 
 
 plink2zarr.add_command(convert_plink)
+
+
+@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 vcf_partition(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)))

bio2zarr/vcf.py

@@ -236,14 +236,14 @@ def scan_vcfs(paths, show_progress, target_num_partitions):
         regions = indexed_vcf.partition_into_regions(num_parts=target_num_partitions)
         for region in regions:
             partitions.append(
-                # Requires cyvcf2>=0.30.27
                 VcfPartition(
                     vcf_path=str(path),
                     region=region,
                 )
             )
-    # TODO figure out if this is safe when we have multiple chrs
-    # in the file
+    # TODO figure out if this is correct. It should be fine because of VCF
+    # sorting, but need to verify there isn't some loophole with BCF, where
+    # contigs are sorted by their index in the list rather than string value.
     partitions.sort(key=lambda x: (x.region.contig, x.region.start))
     vcf_metadata.partitions = partitions
     return vcf_metadata, header
@@ -791,7 +791,7 @@ def convert_partition(
         partition = vcf_metadata.partitions[partition_index]
         vcf = cyvcf2.VCF(partition.vcf_path)
         logger.info(
-            f"Start partition {partition_index} {partition.vcf_path}: {partition.region}"
+            f"Start p{partition_index} {partition.vcf_path}__{partition.region}"
         )
 
         info_fields = []
@@ -835,6 +835,7 @@ def variants():
             encoder_threads=0,
             chunk_size=column_chunk_size,
         ) as tcw:
+            num_records = 0
             for variant in variants():
                 tcw.append("CHROM", variant.CHROM)
                 tcw.append("POS", variant.POS)
@@ -859,8 +860,13 @@ def variants():
                 # is that we get a significant pause at the end of the counter while
                 # all the "small" fields get flushed. Possibly not much to be done about it.
                 core.update_progress(1)
+                num_records += 1
 
-        return tcw.field_summaries
+        logger.info(
+            f"Finish p{partition_index} {partition.vcf_path}__{partition.region}="
+            f"{num_records} records")
+
+        return tcw.field_summaries, num_records
 
     @staticmethod
     def convert(
@@ -896,7 +902,13 @@ def convert(
                     out_path,
                     column_chunk_size=column_chunk_size,
                 )
-            partition_summaries = list(pwm.results_as_completed())
+            num_records = 0
+            partition_summaries = []
+            for partition_summary, partition_num_records in pwm.results_as_completed():
+                num_records += partition_num_records
+                partition_summaries.append(partition_summary)
+
+        assert num_records == pcvcf.num_records
 
         for field in vcf_metadata.fields:
             # Clear the summary to avoid problems when running in debug

bio2zarr/vcf_utils.py

@@ -23,8 +23,6 @@
 )
 
 
-# TODO create a Region dataclass that will sort correctly, and has
-# a str method that does the correct thing
 
 
 def region_filter(variants, region=None):
@@ -62,8 +60,10 @@ class Region:
     def __post_init__(self):
         if self.start is not None:
             self.start = int(self.start)
+            assert self.start > 0
         if self.end is not None:
             self.end = int(self.end)
+            assert self.end > 0
 
     def __str__(self):
         s = f"{self.contig}"

tests/test_cli.py

@@ -122,7 +122,26 @@ def test_convert_plink(self):
             )
 
 
-@pytest.mark.parametrize("cmd", [main.bio2zarr, cli.vcf2zarr, cli.plink2zarr])
+class TestVcfPartition:
+    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.vcf_partition, [path, "-n", "5"], catch_exceptions=False
+        )
+        assert list(result.stdout.splitlines()) == [
+            "20:1-278528",
+            "20:278529-442368",
+            "20:442369-638976",
+            "20:638977-819200",
+            "20:819201-",
+        ]
+
+
+@pytest.mark.parametrize(
+    "cmd", [main.bio2zarr, cli.vcf2zarr, cli.plink2zarr, cli.vcf_partition]
+)
 def test_version(cmd):
     runner = ct.CliRunner(mix_stderr=False)
     result = runner.invoke(cmd, ["--version"], catch_exceptions=False)