add default compressor to config

Author: brokkoli71

src/zarr/core/config.py

@@ -64,6 +64,10 @@ 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",
+            },
         }
     ],
 )

src/zarr/core/metadata/v2.py

@@ -71,6 +71,8 @@ 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)
         compressor_parsed = parse_compressor(compressor)
         order_parsed = parse_indexing_order(order)
         dimension_separator_parsed = parse_separator(dimension_separator)
@@ -238,15 +240,15 @@ def parse_filters(data: object) -> tuple[numcodecs.abc.Codec, ...] | None:
     raise TypeError(msg)
 
 
-def parse_compressor(data: object) -> numcodecs.abc.Codec | None:
+def parse_compressor(data: object) -> numcodecs.abc.Codec:
     """
     Parse a potential compressor.
     """
-    if data is None or isinstance(data, numcodecs.abc.Codec):
+    if isinstance(data, numcodecs.abc.Codec):
         return data
     if isinstance(data, dict):
         return numcodecs.get_codec(data)
-    msg = f"Invalid compressor. Expected None, a numcodecs.abc.Codec, or a dict representation of a numcodecs.abc.Codec. Got {type(data)} instead."
+    msg = f"Invalid compressor. Expected a numcodecs.abc.Codec, or a dict representation of a numcodecs.abc.Codec. Got {type(data)} instead."
     raise ValueError(msg)
 
 
@@ -326,3 +328,16 @@ def _default_fill_value(dtype: np.dtype[Any]) -> Any:
         return ""
     else:
         return dtype.type(0)
+
+
+def _default_compressor(dtype: np.dtype[Any]) -> numcodecs.abc.Codec:
+    """Get the default compressor for a type.
+
+    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():
+        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}")

tests/test_config.py

@@ -63,6 +63,10 @@ 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",
+            },
         }
     ]
     assert config.get("array.order") == "C"

tests/test_v2.py

@@ -10,7 +10,7 @@
 
 import zarr
 import zarr.storage
-from zarr import Array
+from zarr import Array, config
 from zarr.storage import MemoryStore, StorePath
 
 
@@ -96,7 +96,6 @@ async def test_v2_encode_decode(dtype):
     serialized = json.loads(result.to_bytes())
     expected = {
         "chunks": [3],
-        "compressor": None,
         "dtype": f"{dtype}0",
         "fill_value": "WA==",
         "filters": None,
@@ -105,6 +104,7 @@ async def test_v2_encode_decode(dtype):
         "zarr_format": 2,
         "dimension_separator": ".",
     }
+    del serialized["compressor"]
     assert serialized == expected
 
     data = zarr.open_array(store=store, path="foo")[:]
@@ -130,3 +130,21 @@ def test_v2_filters_codecs(filters: Any) -> None:
     arr[:] = array_fixture
     result = arr[:]
     np.testing.assert_array_equal(result, array_fixture)
+
+
+@pytest.mark.parametrize(
+    "dtype_compressor",
+    [["b", "zstd"], ["i", "zstd"], ["f", "zstd"], ["|S1", "vlen-bytes"], ["|U1", "vlen-bytes"]],
+)
+def test_default_compressors(dtype_compressor: Any) -> None:
+    with config.set(
+        {
+            "v2_dtype_kind_to_default_compressor": {
+                "biufcmM": "zstd",
+                "OSUV": "vlen-bytes",
+            },
+        }
+    ):
+        dtype, expected_compressor = dtype_compressor
+        arr = zarr.create(shape=(10,), path="foo", store={}, zarr_format=2, dtype=dtype)
+        assert arr.metadata.compressor.codec_id == expected_compressor