Ensure default compressor is not ignored

Author: tomwhite

sgkit/io/vcf/utils.py

@@ -168,3 +168,32 @@ def temporary_directory(
     finally:
         # Remove the temporary directory on exiting the context manager
         fs.rm(tempdir, recursive=True)
+
+
+def merge_encodings(
+    default_encoding: Dict[str, Dict[str, Any]], overrides: Dict[str, Dict[str, Any]]
+) -> Dict[str, Dict[str, Any]]:
+    """Merge a dictionary of dictionaries specifying encodings with another dictionary of dictionaries of overriding encodings.
+
+    Parameters
+    ----------
+    default_encoding : Dict[str, Dict[str, Any]]
+        The default encoding dictionary.
+    overrides : Dict[str, Dict[str, Any]]
+        A dictionary containing selective overrides.
+
+    Returns
+    -------
+    Dict[str, Dict[str, Any]]
+        The merged encoding dictionary
+    """
+    merged = {}
+    for var, d in default_encoding.items():
+        if var in overrides:
+            merged[var] = {**d, **overrides[var]}
+        else:
+            merged[var] = d
+    for var, d in overrides.items():
+        if var not in merged:
+            merged[var] = d
+    return merged

sgkit/io/vcf/vcf_reader.py

@@ -37,7 +37,13 @@
     STR_MISSING,
 )
 from sgkit.io.vcf import partition_into_regions
-from sgkit.io.vcf.utils import build_url, chunks, temporary_directory, url_filename
+from sgkit.io.vcf.utils import (
+    build_url,
+    chunks,
+    merge_encodings,
+    temporary_directory,
+    url_filename,
+)
 from sgkit.io.vcfzarr_reader import (
     concat_zarrs_optimized,
     vcf_number_to_dimension_and_size,
@@ -556,7 +562,7 @@ def get_chunk_size(dim: Hashable, size: int) -> int:
 
                 # values from function args (encoding) take precedence over default_encoding
                 encoding = encoding or {}
-                merged_encoding = {**default_encoding, **encoding}
+                merged_encoding = merge_encodings(default_encoding, encoding)
 
                 ds.to_zarr(output, mode="w", encoding=merged_encoding)
                 first_variants_chunk = False

sgkit/tests/io/vcf/test_utils.py

@@ -6,7 +6,7 @@
 import pytest
 from callee.strings import StartsWith
 
-from sgkit.io.vcf.utils import build_url, chunks, temporary_directory
+from sgkit.io.vcf.utils import build_url, chunks, merge_encodings, temporary_directory
 from sgkit.io.vcf.vcf_reader import get_region_start
 
 
@@ -118,3 +118,14 @@ def test_chunks(x, n, expected_values):
 )
 def test_get_region_start(region: str, expected: int):
     assert get_region_start(region) == expected
+
+
+def test_merge_encodings():
+    default_encoding = dict(a=dict(a1=1, a2=2), b=dict(b1=5))
+    overrides = dict(a=dict(a1=0, a3=3), c=dict(c1=7))
+    assert merge_encodings(default_encoding, overrides) == dict(
+        a=dict(a1=0, a2=2, a3=3), b=dict(b1=5), c=dict(c1=7)
+    )
+
+    assert merge_encodings(default_encoding, {}) == default_encoding
+    assert merge_encodings({}, overrides) == overrides

sgkit/tests/io/vcf/test_vcf_reader.py

@@ -215,7 +215,7 @@ def test_vcf_to_zarr__compressor_and_filters(shared_datadir, is_path, tmp_path):
     path = path_for_test(shared_datadir, "sample.vcf.gz", is_path)
     output = tmp_path.joinpath("vcf.zarr").as_posix()
 
-    default_compressor = Blosc("zlib", 1, Blosc.NOSHUFFLE)
+    compressor = Blosc("zlib", 1, Blosc.NOSHUFFLE)
     variant_id_compressor = Blosc("zlib", 2, Blosc.NOSHUFFLE)
     encoding = dict(
         variant_id=dict(compressor=variant_id_compressor),
@@ -226,18 +226,25 @@ def test_vcf_to_zarr__compressor_and_filters(shared_datadir, is_path, tmp_path):
         output,
         chunk_length=5,
         chunk_width=2,
-        compressor=default_compressor,
+        compressor=compressor,
         encoding=encoding,
     )
 
     # look at actual Zarr store to check compressor and filters
     z = zarr.open(output)
-    assert z["call_genotype"].compressor == default_compressor
-    assert z["call_genotype"].filters is None
-    assert z["call_genotype_mask"].filters == [PackBits()]
+    assert z["call_genotype"].compressor == compressor
+    assert z["call_genotype"].filters is None  # sgkit default
+    assert z["call_genotype"].chunks == (5, 2, 2)
+    assert z["call_genotype_mask"].compressor == compressor
+    assert z["call_genotype_mask"].filters == [PackBits()]  # sgkit default
+    assert z["call_genotype_mask"].chunks == (5, 2, 2)
 
     assert z["variant_id"].compressor == variant_id_compressor
+    assert z["variant_id"].filters == [VLenUTF8()]  # sgkit default
+    assert z["variant_id"].chunks == (5,)
+    assert z["variant_id_mask"].compressor == compressor
     assert z["variant_id_mask"].filters is None
+    assert z["variant_id_mask"].chunks == (5,)
 
 
 @pytest.mark.parametrize(