Add local-alleles option to CLI

Author: Will-Tyler

bio2zarr/cli.py

@@ -149,6 +149,14 @@ def list_commands(self, ctx):
     help="An approximate bound on overall memory usage (e.g. 10G),",
 )
 
+local_alleles = click.option(
+    "--local-alleles",
+    type=bool,
+    show_default=True,
+    default=True,
+    help="Use local allele fields to reduce the storage requirements of the output.",
+)
+
 
 def setup_logging(verbosity):
     level = "WARNING"
@@ -214,6 +222,7 @@ def show_work_summary(work_summary, json):
 @compressor
 @progress
 @worker_processes
+@local_alleles
 def explode(
     vcfs,
     icf_path,
@@ -223,6 +232,7 @@ def explode(
     compressor,
     progress,
     worker_processes,
+    local_alleles,
 ):
     """
     Convert VCF(s) to intermediate columnar format
@@ -236,6 +246,7 @@ def explode(
         column_chunk_size=column_chunk_size,
         compressor=get_compressor(compressor),
         show_progress=progress,
+        local_alleles=local_alleles,
     )
 
 
@@ -250,6 +261,7 @@ def explode(
 @verbose
 @progress
 @worker_processes
+@local_alleles
 def dexplode_init(
     vcfs,
     icf_path,
@@ -261,6 +273,7 @@ def dexplode_init(
     verbose,
     progress,
     worker_processes,
+    local_alleles,
 ):
     """
     Initial step for distributed conversion of VCF(s) to intermediate columnar format
@@ -277,6 +290,7 @@ def dexplode_init(
         worker_processes=worker_processes,
         compressor=get_compressor(compressor),
         show_progress=progress,
+        local_alleles=local_alleles,
     )
     show_work_summary(work_summary, json)
 

bio2zarr/vcf2zarr/icf.py

@@ -216,7 +216,7 @@ def make_field_def(name, vcf_type, vcf_number):
     return fields
 
 
-def scan_vcf(path, target_num_partitions):
+def scan_vcf(path, target_num_partitions, *, local_alleles):
     with vcf_utils.IndexedVcf(path) as indexed_vcf:
         vcf = indexed_vcf.vcf
         filters = []
@@ -1011,6 +1011,7 @@ def init(
         target_num_partitions=None,
         show_progress=False,
         compressor=None,
+        local_alleles,
     ):
         if self.path.exists():
             raise ValueError(f"ICF path already exists: {self.path}")
@@ -1025,6 +1026,7 @@ def init(
             worker_processes=worker_processes,
             show_progress=show_progress,
             target_num_partitions=target_num_partitions,
+            local_alleles=local_alleles,
         )
         check_field_clobbering(icf_metadata)
         self.metadata = icf_metadata
@@ -1245,6 +1247,7 @@ def explode(
     worker_processes=1,
     show_progress=False,
     compressor=None,
+    local_alleles=True,
 ):
     writer = IntermediateColumnarFormatWriter(icf_path)
     writer.init(
@@ -1255,6 +1258,7 @@ def explode(
         show_progress=show_progress,
         column_chunk_size=column_chunk_size,
         compressor=compressor,
+        local_alleles=local_alleles,
     )
     writer.explode(worker_processes=worker_processes, show_progress=show_progress)
     writer.finalise()
@@ -1270,6 +1274,7 @@ def explode_init(
     worker_processes=1,
     show_progress=False,
     compressor=None,
+    local_alleles=True,
 ):
     writer = IntermediateColumnarFormatWriter(icf_path)
     return writer.init(
@@ -1279,6 +1284,7 @@ def explode_init(
         show_progress=show_progress,
         column_chunk_size=column_chunk_size,
         compressor=compressor,
+        local_alleles=local_alleles,
     )
 
 

tests/test_cli.py

@@ -14,6 +14,7 @@
     compressor=None,
     worker_processes=1,
     show_progress=True,
+    local_alleles=True,
 )
 
 DEFAULT_DEXPLODE_PARTITION_ARGS = dict()
@@ -23,6 +24,7 @@
     column_chunk_size=64,
     compressor=None,
     show_progress=True,
+    local_alleles=True,
 )
 
 DEFAULT_ENCODE_ARGS = dict(
@@ -287,6 +289,24 @@ def test_vcf_explode_missing_and_existing_vcf(self, mocked, tmp_path):
         assert "'no_such_file' does not exist" in result.stderr
         mocked.assert_not_called()
 
+    @pytest.mark.parametrize("local_alleles", [False, True])
+    @mock.patch("bio2zarr.vcf2zarr.explode")
+    def test_vcf_explode_local_alleles(self, mocked, tmp_path, local_alleles):
+        icf_path = tmp_path / "icf"
+        runner = ct.CliRunner(mix_stderr=False)
+        result = runner.invoke(
+            cli.vcf2zarr_main,
+            f"explode {self.vcf_path} {icf_path}"
+            f" --local-alleles {str(local_alleles).lower()}",
+            catch_exceptions=False,
+        )
+        assert result.exit_code == 0
+        assert len(result.stdout) == 0
+        assert len(result.stderr) == 0
+        args = dict(DEFAULT_EXPLODE_ARGS)
+        args["local_alleles"] = local_alleles
+        mocked.assert_called_once_with(str(icf_path), (self.vcf_path,), **args)
+
     @pytest.mark.parametrize(("progress", "flag"), [(True, "-P"), (False, "-Q")])
     @mock.patch("bio2zarr.vcf2zarr.explode_init", return_value=FakeWorkSummary(5))
     def test_vcf_dexplode_init(self, mocked, tmp_path, progress, flag):

tests/test_icf.py

@@ -24,6 +24,11 @@ class TestSmallExample:
 
     @pytest.fixture(scope="class")
     def icf(self, tmp_path_factory):
+        out = tmp_path_factory.mktemp("data") / "example.exploded"
+        return vcf2zarr.explode(out, [self.data_path], local_alleles=False)
+
+    @pytest.fixture(scope="class")
+    def icf_local_alleles(self, tmp_path_factory):
         out = tmp_path_factory.mktemp("data") / "example.exploded"
         return vcf2zarr.explode(out, [self.data_path])
 
@@ -49,6 +54,13 @@ def test_summary_table(self, icf):
         fields = [d["name"] for d in data]
         assert tuple(sorted(fields)) == self.fields
 
+    def test_summary_table_local_allleles(self, icf_local_alleles):
+        data = icf_local_alleles.summary_table()
+        fields = [d["name"] for d in data]
+        fields.sort()
+        expected = tuple(sorted((*self.fields, "FORMAT/LAA")))
+        assert tuple(fields) == expected
+
     def test_inspect(self, icf):
         assert icf.summary_table() == vcf2zarr.inspect(icf.path)