modify _default_compressor to _default_filters_and_compressor

Author: brokkoli71

src/zarr/core/array.py

@@ -492,14 +492,6 @@ async def create(
                 order=order,
             )
         elif zarr_format == 2:
-            if dtype is str or dtype == "str":
-                # another special case: zarr v2 added the vlen-utf8 codec
-                vlen_codec: dict[str, JSON] = {"id": "vlen-utf8"}
-                if filters and not any(x["id"] == "vlen-utf8" for x in filters):
-                    filters = list(filters) + [vlen_codec]
-                else:
-                    filters = [vlen_codec]
-
             if codecs is not None:
                 raise ValueError(
                     "codecs cannot be used for arrays with version 2. Use filters and compressor instead."

src/zarr/core/config.py

@@ -64,9 +64,9 @@ def reset(self) -> None:
             },
             "buffer": "zarr.core.buffer.cpu.Buffer",
             "ndbuffer": "zarr.core.buffer.cpu.NDBuffer",
-            "v2_dtype_kind_to_default_compressor": {
-                "biufcmM": "zstd",
-                "OSUV": "vlen-bytes",
+            "v2_dtype_kind_to_default_filters_and_compressor": {
+                "biufcmM": ["zstd"],
+                "OSUV": ["vlen-utf8"],
             },
         }
     ],

src/zarr/core/metadata/v2.py

@@ -4,7 +4,7 @@
 from collections.abc import Iterable
 from enum import Enum
 from functools import cached_property
-from typing import TYPE_CHECKING, TypedDict, cast
+from typing import TYPE_CHECKING, Any, TypedDict, cast
 
 from zarr.abc.metadata import Metadata
 
@@ -71,8 +71,14 @@ def __init__(
         shape_parsed = parse_shapelike(shape)
         dtype_parsed = parse_dtype(dtype)
         chunks_parsed = parse_shapelike(chunks)
-        if compressor is None:
-            compressor = _default_compressor(dtype_parsed)
+        if not filters and not compressor:
+            filters, compressor = _default_filters_and_compressor(dtype_parsed)
+        if dtype is str or dtype == "str":
+            vlen_codec: dict[str, JSON] = {"id": "vlen-utf8"}
+            if filters and not any(x["id"] == "vlen-utf8" for x in filters):
+                filters = list(filters) + [vlen_codec]
+            else:
+                filters = [vlen_codec]
         compressor_parsed = parse_compressor(compressor)
         order_parsed = parse_indexing_order(order)
         dimension_separator_parsed = parse_separator(dimension_separator)
@@ -240,15 +246,15 @@ def parse_filters(data: object) -> tuple[numcodecs.abc.Codec, ...] | None:
     raise TypeError(msg)
 
 
-def parse_compressor(data: object) -> numcodecs.abc.Codec:
+def parse_compressor(data: object) -> numcodecs.abc.Codec | None:
     """
     Parse a potential compressor.
     """
-    if isinstance(data, numcodecs.abc.Codec):
+    if data is None or isinstance(data, numcodecs.abc.Codec):
         return data
     if isinstance(data, dict):
         return numcodecs.get_codec(data)
-    msg = f"Invalid compressor. Expected a numcodecs.abc.Codec, or a dict representation of a numcodecs.abc.Codec. Got {type(data)} instead."
+    msg = f"Invalid compressor. Expected None, a numcodecs.abc.Codec, or a dict representation of a numcodecs.abc.Codec. Got {type(data)} instead."
     raise ValueError(msg)
 
 
@@ -330,14 +336,18 @@ def _default_fill_value(dtype: np.dtype[Any]) -> Any:
         return dtype.type(0)
 
 
-def _default_compressor(dtype: np.dtype[Any]) -> numcodecs.abc.Codec:
-    """Get the default compressor for a type.
+def _default_filters_and_compressor(
+    dtype: np.dtype[Any],
+) -> tuple[list[dict[str, str]], dict[str, str] | None]:
+    """Get the default filters and compressor for a dtype.
 
     The config contains a mapping from numpy dtype kind to the default compressor.
     https://numpy.org/doc/2.1/reference/generated/numpy.dtype.kind.html
     """
-    dtype_kind_to_default_compressor = config.get("v2_dtype_kind_to_default_compressor")
-    for dtype_kinds, compressor in dtype_kind_to_default_compressor.items():
+    dtype_kind_to_default_compressor = config.get("v2_dtype_kind_to_default_filters_and_compressor")
+    for dtype_kinds, filters_and_compressor in dtype_kind_to_default_compressor.items():
         if dtype.kind in dtype_kinds:
-            return numcodecs.get_codec({"id": compressor})
-    raise ValueError(f"No default compressor found for dtype {dtype} of kind {dtype.kind}")
+            filters = [{"id": f} for f in filters_and_compressor]
+            compressor = None
+            return filters, compressor
+    return [], None

tests/test_config.py

@@ -63,9 +63,9 @@ def test_config_defaults_set() -> None:
                 "vlen-utf8": "zarr.codecs.vlen_utf8.VLenUTF8Codec",
                 "vlen-bytes": "zarr.codecs.vlen_utf8.VLenBytesCodec",
             },
-            "v2_dtype_kind_to_default_compressor": {
-                "biufcmM": "zstd",
-                "OSUV": "vlen-bytes",
+            "v2_dtype_kind_to_default_filters_and_compressor": {
+                "biufcmM": ["zstd"],
+                "OSUV": ["vlen-utf8"],
             },
         }
     ]

tests/test_v2.py

@@ -80,36 +80,43 @@ def test_codec_pipeline() -> None:
 
 @pytest.mark.parametrize("dtype", ["|S", "|V"])
 async def test_v2_encode_decode(dtype):
-    store = zarr.storage.MemoryStore(mode="w")
-    g = zarr.group(store=store, zarr_format=2)
-    g.create_array(
-        name="foo",
-        shape=(3,),
-        chunks=(3,),
-        dtype=dtype,
-        fill_value=b"X",
-    )
-
-    result = await store.get("foo/.zarray", zarr.core.buffer.default_buffer_prototype())
-    assert result is not None
-
-    serialized = json.loads(result.to_bytes())
-    expected = {
-        "chunks": [3],
-        "dtype": f"{dtype}0",
-        "fill_value": "WA==",
-        "filters": None,
-        "order": "C",
-        "shape": [3],
-        "zarr_format": 2,
-        "dimension_separator": ".",
-    }
-    del serialized["compressor"]
-    assert serialized == expected
+    with config.set(
+        {
+            "v2_dtype_kind_to_default_filters_and_compressor": {
+                "OSUV": ["vlen-bytes"],
+            },
+        }
+    ):
+        store = zarr.storage.MemoryStore(mode="w")
+        g = zarr.group(store=store, zarr_format=2)
+        g.create_array(
+            name="foo",
+            shape=(3,),
+            chunks=(3,),
+            dtype=dtype,
+            fill_value=b"X",
+        )
+
+        result = await store.get("foo/.zarray", zarr.core.buffer.default_buffer_prototype())
+        assert result is not None
+
+        serialized = json.loads(result.to_bytes())
+        expected = {
+            "chunks": [3],
+            "compressor": None,
+            "dtype": f"{dtype}0",
+            "fill_value": "WA==",
+            "filters": [{"id": "vlen-bytes"}],
+            "order": "C",
+            "shape": [3],
+            "zarr_format": 2,
+            "dimension_separator": ".",
+        }
+        assert serialized == expected
 
-    data = zarr.open_array(store=store, path="foo")[:]
-    expected = np.full((3,), b"X", dtype=dtype)
-    np.testing.assert_equal(data, expected)
+        data = zarr.open_array(store=store, path="foo")[:]
+        expected = np.full((3,), b"X", dtype=dtype)
+        np.testing.assert_equal(data, expected)
 
 
 @pytest.mark.parametrize("dtype", [str, "str"])
@@ -133,18 +140,19 @@ def test_v2_filters_codecs(filters: Any) -> None:
 
 
 @pytest.mark.parametrize(
-    "dtype_compressor",
-    [["b", "zstd"], ["i", "zstd"], ["f", "zstd"], ["|S1", "vlen-bytes"], ["|U1", "vlen-bytes"]],
+    "dtype_expected",
+    [["b", "zstd"], ["i", "zstd"], ["f", "zstd"], ["|S1", "vlen-utf8"], ["|U1", "vlen-utf8"]],
 )
-def test_default_compressors(dtype_compressor: Any) -> None:
+def test_default_filters_and_compressor(dtype_expected: Any) -> None:
     with config.set(
         {
-            "v2_dtype_kind_to_default_compressor": {
-                "biufcmM": "zstd",
-                "OSUV": "vlen-bytes",
+            "v2_dtype_kind_to_default_filters_and_compressor": {
+                "biufcmM": ["zstd"],
+                "OSUV": ["vlen-utf8"],
             },
         }
     ):
-        dtype, expected_compressor = dtype_compressor
+        dtype, expected = dtype_expected
         arr = zarr.create(shape=(10,), path="foo", store={}, zarr_format=2, dtype=dtype)
-        assert arr.metadata.compressor.codec_id == expected_compressor
+        assert arr.metadata.filters[0].codec_id == expected
+        print(arr.metadata)