Simplify distributed encode

Closes #138
Closes #139

Author: jeromekelleher

CHANGELOG.md

@@ -1,3 +1,8 @@
+# 0.0.7 2024-04-xx
+- Change on-disk format of distributed encode and simplify
+- Check for all partitions nominally completed encoding before doing
+  anything destructive in dencode-finalise
+
 # 0.0.6 2024-04-24
 
 - Only use NOSHUFFLE by default on ``call_genotype`` and bool arrays.

bio2zarr/vcf.py

@@ -1758,28 +1758,42 @@ def load_metadata(self):
     def partition_path(self, partition_index):
         return self.partitions_path / f"p{partition_index}"
 
+    def wip_partition_path(self, partition_index):
+        return self.partitions_path / f"wip_p{partition_index}"
+
     def wip_partition_array_path(self, partition_index, name):
-        return self.partition_path(partition_index) / f"wip_{name}"
+        return self.wip_partition_path(partition_index) / name
 
     def partition_array_path(self, partition_index, name):
         return self.partition_path(partition_index) / name
 
     def encode_partition(self, partition_index):
         self.load_metadata()
-        partition_path = self.partition_path(partition_index)
+        if partition_index < 0 or partition_index >= self.num_partitions:
+            raise ValueError(
+                "Partition index must be in the range 0 <= index < num_partitions"
+            )
+        partition_path = self.wip_partition_path(partition_index)
         partition_path.mkdir(exist_ok=True)
         logger.info(f"Encoding partition {partition_index} to {partition_path}")
 
-        self.encode_alleles_partition(partition_index)
         self.encode_id_partition(partition_index)
         self.encode_filters_partition(partition_index)
         self.encode_contig_partition(partition_index)
+        self.encode_alleles_partition(partition_index)
         for col in self.schema.columns.values():
             if col.vcf_field is not None:
                 self.encode_array_partition(col, partition_index)
         if "call_genotype" in self.schema.columns:
             self.encode_genotypes_partition(partition_index)
 
+        final_path = self.partition_path(partition_index)
+        logger.info(f"Finalising {partition_index} at {final_path}")
+        if final_path.exists():
+            logger.warning("Removing existing partition at {final_path}")
+            shutil.rmtree(final_path)
+        os.rename(partition_path, final_path)
+
     def init_partition_array(self, partition_index, name):
         wip_path = self.wip_partition_array_path(partition_index, name)
         # Create an empty array like the definition
@@ -1791,14 +1805,6 @@ def init_partition_array(self, partition_index, name):
         return array
 
     def finalise_partition_array(self, partition_index, name):
-        wip_path = self.wip_partition_array_path(partition_index, name)
-        final_path = self.partition_array_path(partition_index, name)
-        if final_path.exists():
-            # NEEDS TEST
-            logger.warning(f"Removing existing {final_path}")
-            shutil.rmtree(final_path)
-        # Atomic swap
-        os.rename(wip_path, final_path)
         logger.debug(f"Encoded {name} partition {partition_index}")
 
     def encode_array_partition(self, column, partition_index):
@@ -1942,7 +1948,7 @@ def finalise_array(self, name):
         if final_path.exists():
             # NEEDS TEST
             raise ValueError(f"Array {name} already exists")
-        for partition in range(len(self.metadata.partitions)):
+        for partition in range(self.num_partitions):
             # Move all the files in partition dir to dest dir
             src = self.partition_array_path(partition, name)
             if not src.exists():
@@ -1969,6 +1975,13 @@ def finalise_array(self, name):
     def finalise(self, show_progress=False):
         self.load_metadata()
 
+        missing = []
+        for partition_id in range(self.num_partitions):
+            if not self.partition_path(partition_id).exists():
+                missing.append(partition_id)
+        if len(missing) > 0:
+            raise FileNotFoundError(f"Partitions not encoded: {missing}")
+
         progress_config = core.ProgressConfig(
             total=len(self.schema.columns),
             title="Finalise",
@@ -1986,6 +1999,9 @@ def finalise(self, show_progress=False):
         with core.ParallelWorkManager(0, progress_config) as pwm:
             for name in self.schema.columns:
                 pwm.submit(self.finalise_array, name)
+        logger.debug(f"Removing {self.wip_path}")
+        shutil.rmtree(self.wip_path)
+        logger.info("Consolidating Zarr metadata")
         zarr.consolidate_metadata(self.path)
 
     ######################

tests/test_vcf.py

@@ -5,7 +5,7 @@
 import xarray.testing as xt
 import zarr
 
-from bio2zarr import vcf, vcf_utils
+from bio2zarr import core, vcf, vcf_utils
 
 
 @pytest.fixture(scope="module")
@@ -383,3 +383,117 @@ def test_fields(self, schema, field, description):
     #     ])
     # def test_filters(self, schema, filt, description):
     #     assert schema["filters"][field]["description"] == description
+
+
+class TestVcfZarrWriterExample:
+    arrays = (
+        "variant_contig",
+        "variant_filter",
+        "variant_id",
+        "variant_AA",
+        "variant_AC",
+        "variant_AF",
+        "variant_AN",
+        "variant_DB",
+        "variant_DP",
+        "variant_H2",
+        "variant_NS",
+        "variant_position",
+        "variant_quality",
+        "variant_allele",
+        "call_DP",
+        "call_GQ",
+        "call_genotype",
+        "call_genotype_phased",
+        "call_genotype_mask",
+        "call_HQ",
+    )
+
+    def test_init_paths(self, icf_path, tmp_path):
+        zarr_path = tmp_path / "x.zarr"
+        assert not zarr_path.exists()
+        num_partitions, _ = vcf.encode_init(
+            icf_path, zarr_path, 7, variants_chunk_size=3
+        )
+        assert num_partitions == 3
+        assert zarr_path.exists()
+        wip_path = zarr_path / "wip"
+        assert wip_path.exists()
+        wip_partitions_path = wip_path / "partitions"
+        assert wip_partitions_path.exists()
+        wip_arrays_path = wip_path / "arrays"
+        assert wip_arrays_path.exists()
+        for name in self.arrays:
+            array_path = wip_arrays_path / name
+            assert array_path.exists()
+        with open(wip_path / "metadata.json") as f:
+            d = json.loads(f.read())
+            # Basic test
+            assert len(d["partitions"]) == 3
+
+    def test_finalise_paths(self, icf_path, tmp_path):
+        zarr_path = tmp_path / "x.zarr"
+        assert not zarr_path.exists()
+        num_partitions, _ = vcf.encode_init(
+            icf_path, zarr_path, 7, variants_chunk_size=3
+        )
+        wip_path = zarr_path / "wip"
+        assert wip_path.exists()
+        for j in range(num_partitions):
+            vcf.encode_partition(zarr_path, j)
+            assert (wip_path / "partitions" / f"p{j}").exists()
+        vcf.encode_finalise(zarr_path)
+        assert zarr_path.exists()
+        assert not wip_path.exists()
+
+    def test_finalise_no_partitions_fails(self, icf_path, tmp_path):
+        zarr_path = tmp_path / "x.zarr"
+        vcf.encode_init(icf_path, zarr_path, 3, variants_chunk_size=3)
+        with pytest.raises(
+            FileNotFoundError, match="Partitions not encoded: \\[0, 1, 2\\]"
+        ):
+            vcf.encode_finalise(zarr_path)
+
+    @pytest.mark.parametrize("partition", [0, 1, 2])
+    def test_finalise_missing_partition_fails(self, icf_path, tmp_path, partition):
+        zarr_path = tmp_path / "x.zarr"
+        vcf.encode_init(icf_path, zarr_path, 3, variants_chunk_size=3)
+        for j in range(3):
+            if j != partition:
+                vcf.encode_partition(zarr_path, j)
+        with pytest.raises(
+            FileNotFoundError, match=f"Partitions not encoded: \\[{partition}\\]"
+        ):
+            vcf.encode_finalise(zarr_path)
+
+    @pytest.mark.parametrize("partition", [0, 1, 2])
+    def test_encode_partition(self, icf_path, tmp_path, partition):
+        zarr_path = tmp_path / "x.zarr"
+        vcf.encode_init(icf_path, zarr_path, 3, variants_chunk_size=3)
+        partition_path = zarr_path / "wip" / "partitions" / f"p{partition}"
+        assert not partition_path.exists()
+        vcf.encode_partition(zarr_path, partition)
+        assert partition_path.exists()
+
+    def test_double_encode_partition(self, icf_path, tmp_path, caplog):
+        partition = 1
+        zarr_path = tmp_path / "x.zarr"
+        vcf.encode_init(icf_path, zarr_path, 3, variants_chunk_size=3)
+        partition_path = zarr_path / "wip" / "partitions" / f"p{partition}"
+        assert not partition_path.exists()
+        vcf.encode_partition(zarr_path, partition)
+        assert partition_path.exists()
+        size = core.du(partition_path)
+        assert size > 0
+        with caplog.at_level("WARNING"):
+            vcf.encode_partition(zarr_path, partition)
+        assert "Removing existing partition at" in caplog.text
+        assert partition_path.exists()
+        assert core.du(partition_path) == size
+
+    @pytest.mark.parametrize("partition", [-1, 3, 100])
+    def test_encode_partition_out_of_range(self, icf_path, tmp_path, partition):
+        zarr_path = tmp_path / "x.zarr"
+        vcf.encode_init(icf_path, zarr_path, 3, variants_chunk_size=3)
+        with pytest.raises(ValueError, match="Partition index must be in the range"):
+            vcf.encode_partition(zarr_path, partition)