Merge pull request #116 from jeromekelleher/schema-bug

Fix ColumnSpec initialisation

Author: jeromekelleher

CHANGELOG.md

@@ -1,3 +1,7 @@
+# 0.0.5 2024-04-XX
+
+- Fix bug in schema handling (compressor settings ignored)
+
 # 0.0.4 2024-04-08
 
 - Fix bug in --max-memory handling, and argument to a string like 10G

bio2zarr/vcf.py

@@ -890,7 +890,6 @@ def num_columns(self):
         return len(self.columns)
 
 
-
 def mkdir_with_progress(path):
     logger.debug(f"mkdir f{path}")
     # NOTE we may have race-conditions here, I'm not sure. Hopefully allowing
@@ -1226,20 +1225,25 @@ class ZarrColumnSpec:
     dtype: str
     shape: tuple
     chunks: tuple
-    dimensions: list
+    dimensions: tuple
     description: str
     vcf_field: str
-    compressor: dict = None
-    filters: list = None
-    # TODO add filters
+    compressor: dict
+    filters: list
 
     def __post_init__(self):
+        # Ensure these are tuples for ease of comparison and consistency
         self.shape = tuple(self.shape)
         self.chunks = tuple(self.chunks)
         self.dimensions = tuple(self.dimensions)
-        self.compressor = DEFAULT_ZARR_COMPRESSOR.get_config()
-        self.filters = []
-        self._choose_compressor_settings()
+
+    @staticmethod
+    def new(**kwargs):
+        spec = ZarrColumnSpec(
+            **kwargs, compressor=DEFAULT_ZARR_COMPRESSOR.get_config(), filters=[]
+        )
+        spec._choose_compressor_settings()
+        return spec
 
     def _choose_compressor_settings(self):
         """
@@ -1315,7 +1319,7 @@ def generate(icf, variants_chunk_size=None, samples_chunk_size=None):
         def fixed_field_spec(
             name, dtype, vcf_field=None, shape=(m,), dimensions=("variants",)
         ):
-            return ZarrColumnSpec(
+            return ZarrColumnSpec.new(
                 vcf_field=vcf_field,
                 name=name,
                 dtype=dtype,
@@ -1399,7 +1403,7 @@ def fixed_field_spec(
                 else:
                     dimensions.append(f"{field.category}_{field.name}_dim")
             variable_name = prefix + field.name
-            colspec = ZarrColumnSpec(
+            colspec = ZarrColumnSpec.new(
                 vcf_field=field.full_name,
                 name=variable_name,
                 dtype=field.smallest_dtype(),
@@ -1417,7 +1421,7 @@ def fixed_field_spec(
             dimensions = ["variants", "samples"]
 
             colspecs.append(
-                ZarrColumnSpec(
+                ZarrColumnSpec.new(
                     vcf_field=None,
                     name="call_genotype_phased",
                     dtype="bool",
@@ -1430,7 +1434,7 @@ def fixed_field_spec(
             shape += [ploidy]
             dimensions += ["ploidy"]
             colspecs.append(
-                ZarrColumnSpec(
+                ZarrColumnSpec.new(
                     vcf_field=None,
                     name="call_genotype",
                     dtype=gt_field.smallest_dtype(),
@@ -1441,7 +1445,7 @@ def fixed_field_spec(
                 )
             )
             colspecs.append(
-                ZarrColumnSpec(
+                ZarrColumnSpec.new(
                     vcf_field=None,
                     name="call_genotype_mask",
                     dtype="bool",
@@ -1523,7 +1527,9 @@ def __init__(self, path, icf, schema, dimension_separator=None):
         self.schema = schema
         # Default to using nested directories following the Zarr v3 default.
         # This seems to require version 2.17+ to work properly
-        self.dimension_separator = "/" if dimension_separator is None else dimension_separator
+        self.dimension_separator = (
+            "/" if dimension_separator is None else dimension_separator
+        )
         store = zarr.DirectoryStore(self.path)
         self.root = zarr.group(store=store)
 

tests/test_vcf.py

@@ -3,6 +3,7 @@
 import pytest
 import xarray.testing as xt
 import sgkit as sg
+import zarr
 
 from bio2zarr import vcf
 
@@ -100,7 +101,6 @@ def test_exploded_metadata_mismatch(self, tmpdir, icf_path, version):
 
 
 class TestEncodeDimensionSeparator:
-
     @pytest.mark.parametrize("dimension_separator", [None, "/"])
     def test_directories(self, tmp_path, icf_path, dimension_separator):
         zarr_path = tmp_path / "zarr"
@@ -122,6 +122,77 @@ def test_bad_value(self, tmp_path, icf_path, dimension_separator):
             vcf.encode(icf_path, zarr_path, dimension_separator=dimension_separator)
 
 
+class TestSchemaJsonRoundTrip:
+    def assert_json_round_trip(self, schema):
+        schema2 = vcf.VcfZarrSchema.fromjson(schema.asjson())
+        assert schema == schema2
+
+    def test_generated_no_changes(self, icf_path):
+        icf = vcf.IntermediateColumnarFormat(icf_path)
+        self.assert_json_round_trip(vcf.VcfZarrSchema.generate(icf))
+
+    def test_generated_no_columns(self, icf_path):
+        icf = vcf.IntermediateColumnarFormat(icf_path)
+        schema = vcf.VcfZarrSchema.generate(icf)
+        schema.columns.clear()
+        self.assert_json_round_trip(schema)
+
+    def test_generated_no_samples(self, icf_path):
+        icf = vcf.IntermediateColumnarFormat(icf_path)
+        schema = vcf.VcfZarrSchema.generate(icf)
+        schema.sample_id.clear()
+        self.assert_json_round_trip(schema)
+
+    def test_generated_change_dtype(self, icf_path):
+        icf = vcf.IntermediateColumnarFormat(icf_path)
+        schema = vcf.VcfZarrSchema.generate(icf)
+        schema.columns["variant_position"].dtype = "i8"
+        self.assert_json_round_trip(schema)
+
+    def test_generated_change_compressor(self, icf_path):
+        icf = vcf.IntermediateColumnarFormat(icf_path)
+        schema = vcf.VcfZarrSchema.generate(icf)
+        schema.columns["variant_position"].compressor = {"cname": "FAKE"}
+        self.assert_json_round_trip(schema)
+
+
+class TestSchemaEncode:
+    @pytest.mark.parametrize(
+        ["cname", "clevel", "shuffle"], [("lz4", 1, 0), ("zlib", 7, 1), ("zstd", 4, 2)]
+    )
+    def test_codec(self, tmp_path, icf_path, cname, clevel, shuffle):
+        zarr_path = tmp_path / "zarr"
+        icf = vcf.IntermediateColumnarFormat(icf_path)
+        schema = vcf.VcfZarrSchema.generate(icf)
+        for var in schema.columns.values():
+            var.compressor["cname"] = cname
+            var.compressor["clevel"] = clevel
+            var.compressor["shuffle"] = shuffle
+        schema_path = tmp_path / "schema"
+        with open(schema_path, "w") as f:
+            f.write(schema.asjson())
+        vcf.encode(icf_path, zarr_path, schema_path=schema_path)
+        root = zarr.open(zarr_path)
+        for var in schema.columns.values():
+            a = root[var.name]
+            assert a.compressor.cname == cname
+            assert a.compressor.clevel == clevel
+            assert a.compressor.shuffle == shuffle
+
+    @pytest.mark.parametrize("dtype", ["i4", "i8"])
+    def test_genotype_dtype(self, tmp_path, icf_path, dtype):
+        zarr_path = tmp_path / "zarr"
+        icf = vcf.IntermediateColumnarFormat(icf_path)
+        schema = vcf.VcfZarrSchema.generate(icf)
+        schema.columns["call_genotype"].dtype = dtype
+        schema_path = tmp_path / "schema"
+        with open(schema_path, "w") as f:
+            f.write(schema.asjson())
+        vcf.encode(icf_path, zarr_path, schema_path=schema_path)
+        root = zarr.open(zarr_path)
+        assert root["call_genotype"].dtype == dtype
+
+
 class TestDefaultSchema:
     def test_format_version(self, schema):
         assert schema["format_version"] == vcf.ZARR_SCHEMA_FORMAT_VERSION