Tidy up vcfpartition interface

Author: jeromekelleher

bio2zarr/cli.py

@@ -518,10 +518,42 @@ def plink2zarr():
 @click.command
 @version
 @click.argument("vcf_path", type=click.Path())
-@click.option("-n", "--num-parts", type=int, default=None)
-# @click.option("-s", "--part-size", type=int, default=None)
-def vcfpartition(vcf_path, num_parts):
+@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)
+    regions = indexed_vcf.partition_into_regions(
+        num_parts=num_parts, target_part_size=part_size
+    )
     for region in regions:
-        click.echo(region)
+        click.echo(f"{region}\t{vcf_path}")

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.
@@ -491,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]