Initial CLI skeleton for dexplode

Author: jeromekelleher

bio2zarr/cli.py

@@ -39,6 +39,14 @@ def list_commands(self, ctx):
     "zarr_path", type=click.Path(file_okay=False, dir_okay=True)
 )
 
+zarr_path = click.argument(
+    "zarr_path", type=click.Path(exists=True, file_okay=False, dir_okay=True)
+)
+
+num_partitions = click.argument("num_partitions", type=click.IntRange(min=1))
+
+partition = click.argument("partition", type=click.IntRange(min=0))
+
 verbose = click.option("-v", "--verbose", count=True, help="Increase verbosity")
 
 force = click.option(
@@ -92,6 +100,27 @@ def list_commands(self, ctx):
     help="Chunk size in the samples dimension",
 )
 
+schema = click.option("-s", "--schema", default=None, type=click.Path(exists=True))
+
+max_variant_chunks = click.option(
+    "-V",
+    "--max-variant-chunks",
+    type=int,
+    default=None,
+    help=(
+        "Truncate the output in the variants dimension to have "
+        "this number of chunks. Mainly intended to help with "
+        "schema tuning."
+    ),
+)
+
+max_memory = click.option(
+    "-M",
+    "--max-memory",
+    default=None,
+    help="An approximate bound on overall memory usage (e.g. 10G),",
+)
+
 
 def setup_logging(verbosity):
     level = "WARNING"
@@ -158,7 +187,7 @@ def explode(
 @click.command
 @vcfs
 @new_icf_path
-@click.argument("num_partitions", type=click.IntRange(min=1))
+@num_partitions
 @force
 @column_chunk_size
 @compressor
@@ -194,7 +223,7 @@ def dexplode_init(
 
 @click.command
 @icf_path
-@click.argument("partition", type=click.IntRange(min=0))
+@partition
 @verbose
 def dexplode_partition(icf_path, partition, verbose):
     """
@@ -207,14 +236,14 @@ def dexplode_partition(icf_path, partition, verbose):
 
 
 @click.command
-@click.argument("path", type=click.Path(), required=True)
+@icf_path
 @verbose
-def dexplode_finalise(path, verbose):
+def dexplode_finalise(icf_path, verbose):
     """
     Final step for distributed conversion of VCF(s) to intermediate columnar format.
     """
     setup_logging(verbose)
-    vcf.explode_finalise(path)
+    vcf.explode_finalise(icf_path)
 
 
 @click.command
@@ -244,26 +273,11 @@ def mkschema(icf_path):
 @new_zarr_path
 @force
 @verbose
-@click.option("-s", "--schema", default=None, type=click.Path(exists=True))
+@schema
 @variants_chunk_size
 @samples_chunk_size
-@click.option(
-    "-V",
-    "--max-variant-chunks",
-    type=int,
-    default=None,
-    help=(
-        "Truncate the output in the variants dimension to have "
-        "this number of chunks. Mainly intended to help with "
-        "schema tuning."
-    ),
-)
-@click.option(
-    "-M",
-    "--max-memory",
-    default=None,
-    help="An approximate bound on overall memory usage (e.g. 10G),",
-)
+@max_variant_chunks
+@max_memory
 @worker_processes
 def encode(
     icf_path,
@@ -295,6 +309,68 @@ def encode(
     )
 
 
+@click.command
+@icf_path
+@new_zarr_path
+@num_partitions
+@force
+@schema
+@variants_chunk_size
+@samples_chunk_size
+@max_variant_chunks
+@verbose
+def dencode_init(
+    icf_path,
+    zarr_path,
+    num_partitions,
+    force,
+    schema,
+    variants_chunk_size,
+    samples_chunk_size,
+    max_variant_chunks,
+    verbose,
+):
+    """
+    TODO DOCUMENT
+    """
+    setup_logging(verbose)
+    check_overwrite_dir(zarr_path, force)
+    num_partitions = vcf.encode_init(
+        icf_path,
+        zarr_path,
+        target_num_partitions=num_partitions,
+        schema_path=schema,
+        variants_chunk_size=variants_chunk_size,
+        samples_chunk_size=samples_chunk_size,
+        max_v_chunks=max_variant_chunks,
+        show_progress=True,
+    )
+    click.echo(num_partitions)
+
+
+@click.command
+@zarr_path
+@partition
+@verbose
+def dencode_partition(zarr_path, partition, verbose):
+    """
+    TODO DOCUMENT
+    """
+    setup_logging(verbose)
+    vcf.encode_partition(zarr_path, partition, show_progress=False)
+
+
+@click.command
+@zarr_path
+@verbose
+def dencode_finalise(zarr_path, verbose):
+    """
+    TODO DOCUMENT
+    """
+    setup_logging(verbose)
+    vcf.encode_finalise(zarr_path)
+
+
 @click.command(name="convert")
 @vcfs
 @new_zarr_path
@@ -382,6 +458,9 @@ def vcf2zarr():
 vcf2zarr.add_command(dexplode_init)
 vcf2zarr.add_command(dexplode_partition)
 vcf2zarr.add_command(dexplode_finalise)
+vcf2zarr.add_command(dencode_init)
+vcf2zarr.add_command(dencode_partition)
+vcf2zarr.add_command(dencode_finalise)
 
 
 @click.command(name="convert")

bio2zarr/vcf.py

@@ -1947,6 +1947,65 @@ def encode(
     vzw.finalise()
 
 
+def encode_init(
+    icf_path,
+    zarr_path,
+    target_num_partitions,
+    *,
+    schema_path=None,
+    variants_chunk_size=None,
+    samples_chunk_size=None,
+    max_v_chunks=None,
+    dimension_separator=None,
+    max_memory=None,
+    worker_processes=1,
+    show_progress=False,
+):
+    pass
+    # icf = IntermediateColumnarFormat(icf_path)
+    # if schema_path is None:
+    #     schema = VcfZarrSchema.generate(
+    #         icf,
+    #         variants_chunk_size=variants_chunk_size,
+    #         samples_chunk_size=samples_chunk_size,
+    #     )
+    # else:
+    #     logger.info(f"Reading schema from {schema_path}")
+    #     if variants_chunk_size is not None or samples_chunk_size is not None:
+    #         raise ValueError(
+    #             "Cannot specify schema along with chunk sizes"
+    #         )  # NEEDS TEST
+    #     with open(schema_path) as f:
+    #         schema = VcfZarrSchema.fromjson(f.read())
+    # zarr_path = pathlib.Path(zarr_path)
+    # if zarr_path.exists():
+    #     logger.warning(f"Deleting existing {zarr_path}")
+
+    #     shutil.rmtree(zarr_path)
+    # vzw = VcfZarrWriter(zarr_path, icf, schema,
+    # dimension_separator=dimension_separator)
+    # vzw.init()
+    # vzw.encode(
+    #     max_v_chunks=max_v_chunks,
+    #     worker_processes=worker_processes,
+    #     max_memory=max_memory,
+    #     show_progress=show_progress,
+    # )
+    # vzw.finalise()
+
+
+def encode_partition(zarr_path, partition, *, show_progress=False, worker_processes=1):
+    writer = VcfZarrWriter(zarr_path)
+    writer.encode_partition(
+        partition, show_progress=show_progress, worker_processes=worker_processes
+    )
+
+
+def encode_finalise(zarr_path):
+    writer = VcfZarrWriter(zarr_path)
+    writer.finalise()
+
+
 def convert(
     vcfs,
     out_path,

tests/test_cli.py

@@ -22,6 +22,7 @@
     compressor=None,
     show_progress=True,
 )
+
 DEFAULT_ENCODE_ARGS = dict(
     schema_path=None,
     variants_chunk_size=None,
@@ -32,6 +33,14 @@
     show_progress=True,
 )
 
+DEFAULT_DENCODE_INIT_ARGS = dict(
+    schema_path=None,
+    variants_chunk_size=None,
+    samples_chunk_size=None,
+    max_v_chunks=None,
+    show_progress=True,
+)
+
 
 class TestWithMocks:
     vcf_path = "tests/data/vcf/sample.vcf.gz"
@@ -385,6 +394,55 @@ def test_encode(self, mocked, tmp_path):
             **DEFAULT_ENCODE_ARGS,
         )
 
+    @mock.patch("bio2zarr.vcf.encode_init", return_value=10)
+    def test_dencode(self, mocked, tmp_path):
+        icf_path = tmp_path / "icf"
+        icf_path.mkdir()
+        zarr_path = tmp_path / "zarr"
+        runner = ct.CliRunner(mix_stderr=False)
+        result = runner.invoke(
+            cli.vcf2zarr,
+            f"dencode-init {icf_path} {zarr_path} 10",
+            catch_exceptions=False,
+        )
+        assert result.exit_code == 0
+        assert result.stdout == "10\n"
+        assert len(result.stderr) == 0
+        mocked.assert_called_once_with(
+            str(icf_path),
+            str(zarr_path),
+            target_num_partitions=10,
+            **DEFAULT_DENCODE_INIT_ARGS,
+        )
+
+    @mock.patch("bio2zarr.vcf.encode_partition")
+    def test_vcf_dencode_partition(self, mocked, tmp_path):
+        runner = ct.CliRunner(mix_stderr=False)
+        zarr_path = tmp_path / "zarr"
+        zarr_path.mkdir()
+        result = runner.invoke(
+            cli.vcf2zarr,
+            f"dencode-partition {zarr_path} 1",
+            catch_exceptions=False,
+        )
+        assert result.exit_code == 0
+        assert len(result.stdout) == 0
+        assert len(result.stderr) == 0
+        mocked.assert_called_once_with(
+            str(zarr_path), 1, **DEFAULT_DEXPLODE_PARTITION_ARGS
+        )
+
+    @mock.patch("bio2zarr.vcf.encode_finalise")
+    def test_vcf_dencode_finalise(self, mocked, tmp_path):
+        runner = ct.CliRunner(mix_stderr=False)
+        result = runner.invoke(
+            cli.vcf2zarr, f"dencode-finalise {tmp_path}", catch_exceptions=False
+        )
+        assert result.exit_code == 0
+        assert len(result.stdout) == 0
+        assert len(result.stderr) == 0
+        mocked.assert_called_once_with(str(tmp_path))
+
     @mock.patch("bio2zarr.vcf.convert")
     def test_convert_vcf(self, mocked):
         runner = ct.CliRunner(mix_stderr=False)