add CodecChain dataclass and sync codec tests

Introduces CodecChain, a frozen dataclass that chains array-array,
array-bytes, and bytes-bytes codecs with synchronous encode/decode
methods. Pure compute only -- no IO, no threading, no batching.

Also adds sync roundtrip tests for individual codecs (blosc, gzip,
zstd, crc32c, bytes, transpose, vlen) and CodecChain integration tests.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: d-v-b

src/zarr/core/codec_pipeline.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from itertools import islice, pairwise
-from typing import TYPE_CHECKING, Any, TypeVar
+from typing import TYPE_CHECKING, Any, TypeVar, cast
 from warnings import warn
 
 from zarr.abc.codec import (
@@ -13,6 +13,7 @@
     BytesBytesCodec,
     Codec,
     CodecPipeline,
+    SupportsSyncCodec,
 )
 from zarr.core.common import concurrent_map
 from zarr.core.config import config
@@ -68,6 +69,134 @@ def fill_value_or_default(chunk_spec: ArraySpec) -> Any:
         return fill_value
 
 
+@dataclass(frozen=True)
+class CodecChain:
+    """Lightweight codec chain: array-array -> array-bytes -> bytes-bytes.
+
+    Pure compute only -- no IO methods, no threading, no batching.
+    """
+
+    array_array_codecs: tuple[ArrayArrayCodec, ...]
+    array_bytes_codec: ArrayBytesCodec
+    bytes_bytes_codecs: tuple[BytesBytesCodec, ...]
+
+    _all_sync: bool = field(default=False, init=False, repr=False, compare=False)
+
+    def __post_init__(self) -> None:
+        object.__setattr__(
+            self,
+            "_all_sync",
+            all(isinstance(c, SupportsSyncCodec) for c in self),
+        )
+
+    def __iter__(self) -> Iterator[Codec]:
+        yield from self.array_array_codecs
+        yield self.array_bytes_codec
+        yield from self.bytes_bytes_codecs
+
+    @classmethod
+    def from_codecs(cls, codecs: Iterable[Codec]) -> CodecChain:
+        aa, ab, bb = codecs_from_list(list(codecs))
+        return cls(array_array_codecs=aa, array_bytes_codec=ab, bytes_bytes_codecs=bb)
+
+    def resolve_metadata_chain(
+        self, chunk_spec: ArraySpec
+    ) -> tuple[
+        list[tuple[ArrayArrayCodec, ArraySpec]],
+        tuple[ArrayBytesCodec, ArraySpec],
+        list[tuple[BytesBytesCodec, ArraySpec]],
+    ]:
+        """Resolve metadata through the codec chain for a single chunk_spec."""
+        aa_codecs_with_spec: list[tuple[ArrayArrayCodec, ArraySpec]] = []
+        spec = chunk_spec
+        for aa_codec in self.array_array_codecs:
+            aa_codecs_with_spec.append((aa_codec, spec))
+            spec = aa_codec.resolve_metadata(spec)
+
+        ab_codec_with_spec = (self.array_bytes_codec, spec)
+        spec = self.array_bytes_codec.resolve_metadata(spec)
+
+        bb_codecs_with_spec: list[tuple[BytesBytesCodec, ArraySpec]] = []
+        for bb_codec in self.bytes_bytes_codecs:
+            bb_codecs_with_spec.append((bb_codec, spec))
+            spec = bb_codec.resolve_metadata(spec)
+
+        return (aa_codecs_with_spec, ab_codec_with_spec, bb_codecs_with_spec)
+
+    def decode_chunk(
+        self,
+        chunk_bytes: Buffer,
+        chunk_spec: ArraySpec,
+        aa_chain: Iterable[tuple[ArrayArrayCodec, ArraySpec]] | None = None,
+        ab_pair: tuple[ArrayBytesCodec, ArraySpec] | None = None,
+        bb_chain: Iterable[tuple[BytesBytesCodec, ArraySpec]] | None = None,
+    ) -> NDBuffer:
+        """Decode a single chunk through the full codec chain, synchronously.
+
+        Pure compute -- no IO. Only callable when all codecs support sync.
+
+        The optional ``aa_chain``, ``ab_pair``, ``bb_chain`` parameters allow
+        pre-resolved metadata to be reused across many chunks with the same spec.
+        If not provided, ``resolve_metadata_chain`` is called internally.
+        """
+        if aa_chain is None or ab_pair is None or bb_chain is None:
+            aa_chain, ab_pair, bb_chain = self.resolve_metadata_chain(chunk_spec)
+
+        bb_out: Any = chunk_bytes
+        for bb_codec, spec in reversed(list(bb_chain)):
+            bb_out = cast("SupportsSyncCodec", bb_codec)._decode_sync(bb_out, spec)
+
+        ab_codec, ab_spec = ab_pair
+        ab_out: Any = cast("SupportsSyncCodec", ab_codec)._decode_sync(bb_out, ab_spec)
+
+        for aa_codec, spec in reversed(list(aa_chain)):
+            ab_out = cast("SupportsSyncCodec", aa_codec)._decode_sync(ab_out, spec)
+
+        return ab_out  # type: ignore[no-any-return]
+
+    def encode_chunk(
+        self,
+        chunk_array: NDBuffer,
+        chunk_spec: ArraySpec,
+    ) -> Buffer | None:
+        """Encode a single chunk through the full codec chain, synchronously.
+
+        Pure compute -- no IO. Only callable when all codecs support sync.
+        """
+        spec = chunk_spec
+        aa_out: Any = chunk_array
+
+        for aa_codec in self.array_array_codecs:
+            if aa_out is None:
+                return None
+            aa_out = cast("SupportsSyncCodec", aa_codec)._encode_sync(aa_out, spec)
+            spec = aa_codec.resolve_metadata(spec)
+
+        if aa_out is None:
+            return None
+        bb_out: Any = cast("SupportsSyncCodec", self.array_bytes_codec)._encode_sync(aa_out, spec)
+        spec = self.array_bytes_codec.resolve_metadata(spec)
+
+        for bb_codec in self.bytes_bytes_codecs:
+            if bb_out is None:
+                return None
+            bb_out = cast("SupportsSyncCodec", bb_codec)._encode_sync(bb_out, spec)
+            spec = bb_codec.resolve_metadata(spec)
+
+        return bb_out  # type: ignore[no-any-return]
+
+    def compute_encoded_size(self, byte_length: int, array_spec: ArraySpec) -> int:
+        for codec in self:
+            byte_length = codec.compute_encoded_size(byte_length, array_spec)
+            array_spec = codec.resolve_metadata(array_spec)
+        return byte_length
+
+    def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec:
+        for codec in self:
+            chunk_spec = codec.resolve_metadata(chunk_spec)
+        return chunk_spec
+
+
 @dataclass(frozen=True)
 class BatchedCodecPipeline(CodecPipeline):
     """Default codec pipeline.

tests/test_codecs/test_blosc.py

@@ -6,11 +6,12 @@
 from packaging.version import Version
 
 import zarr
+from zarr.abc.codec import SupportsSyncCodec
 from zarr.codecs import BloscCodec
 from zarr.codecs.blosc import BloscShuffle, Shuffle
-from zarr.core.array_spec import ArraySpec
+from zarr.core.array_spec import ArrayConfig, ArraySpec
 from zarr.core.buffer import default_buffer_prototype
-from zarr.core.dtype import UInt16
+from zarr.core.dtype import UInt16, get_data_type_from_native_dtype
 from zarr.storage import MemoryStore, StorePath
 
 
@@ -110,3 +111,27 @@ async def test_typesize() -> None:
     else:
         expected_size = 10216
     assert size == expected_size, msg
+
+
+def test_blosc_codec_supports_sync() -> None:
+    assert isinstance(BloscCodec(), SupportsSyncCodec)
+
+
+def test_blosc_codec_sync_roundtrip() -> None:
+    codec = BloscCodec(typesize=8)
+    arr = np.arange(100, dtype="float64")
+    zdtype = get_data_type_from_native_dtype(arr.dtype)
+    spec = ArraySpec(
+        shape=arr.shape,
+        dtype=zdtype,
+        fill_value=zdtype.cast_scalar(0),
+        config=ArrayConfig(order="C", write_empty_chunks=True),
+        prototype=default_buffer_prototype(),
+    )
+    buf = default_buffer_prototype().buffer.from_array_like(arr.view("B"))
+
+    encoded = codec._encode_sync(buf, spec)
+    assert encoded is not None
+    decoded = codec._decode_sync(encoded, spec)
+    result = np.frombuffer(decoded.as_numpy_array(), dtype="float64")
+    np.testing.assert_array_equal(arr, result)

tests/test_codecs/test_crc32c.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+import numpy as np
+
+from zarr.abc.codec import SupportsSyncCodec
+from zarr.codecs.crc32c_ import Crc32cCodec
+from zarr.core.array_spec import ArrayConfig, ArraySpec
+from zarr.core.buffer import default_buffer_prototype
+from zarr.core.dtype import get_data_type_from_native_dtype
+
+
+def test_crc32c_codec_supports_sync() -> None:
+    assert isinstance(Crc32cCodec(), SupportsSyncCodec)
+
+
+def test_crc32c_codec_sync_roundtrip() -> None:
+    codec = Crc32cCodec()
+    arr = np.arange(100, dtype="float64")
+    zdtype = get_data_type_from_native_dtype(arr.dtype)
+    spec = ArraySpec(
+        shape=arr.shape,
+        dtype=zdtype,
+        fill_value=zdtype.cast_scalar(0),
+        config=ArrayConfig(order="C", write_empty_chunks=True),
+        prototype=default_buffer_prototype(),
+    )
+    buf = default_buffer_prototype().buffer.from_array_like(arr.view("B"))
+
+    encoded = codec._encode_sync(buf, spec)
+    assert encoded is not None
+    decoded = codec._decode_sync(encoded, spec)
+    result = np.frombuffer(decoded.as_numpy_array(), dtype="float64")
+    np.testing.assert_array_equal(arr, result)

tests/test_codecs/test_endian.py

@@ -4,8 +4,12 @@
 import pytest
 
 import zarr
+from zarr.abc.codec import SupportsSyncCodec
 from zarr.abc.store import Store
 from zarr.codecs import BytesCodec
+from zarr.core.array_spec import ArrayConfig, ArraySpec
+from zarr.core.buffer import NDBuffer, default_buffer_prototype
+from zarr.core.dtype import get_data_type_from_native_dtype
 from zarr.storage import StorePath
 
 from .test_codecs import _AsyncArrayProxy
@@ -33,6 +37,31 @@ async def test_endian(store: Store, endian: Literal["big", "little"]) -> None:
     assert np.array_equal(data, readback_data)
 
 
+def test_bytes_codec_supports_sync() -> None:
+    assert isinstance(BytesCodec(), SupportsSyncCodec)
+
+
+def test_bytes_codec_sync_roundtrip() -> None:
+    codec = BytesCodec()
+    arr = np.arange(100, dtype="float64")
+    zdtype = get_data_type_from_native_dtype(arr.dtype)
+    spec = ArraySpec(
+        shape=arr.shape,
+        dtype=zdtype,
+        fill_value=zdtype.cast_scalar(0),
+        config=ArrayConfig(order="C", write_empty_chunks=True),
+        prototype=default_buffer_prototype(),
+    )
+    nd_buf: NDBuffer = default_buffer_prototype().nd_buffer.from_numpy_array(arr)
+
+    codec = codec.evolve_from_array_spec(spec)
+
+    encoded = codec._encode_sync(nd_buf, spec)
+    assert encoded is not None
+    decoded = codec._decode_sync(encoded, spec)
+    np.testing.assert_array_equal(arr, decoded.as_numpy_array())
+
+
 @pytest.mark.filterwarnings("ignore:The endianness of the requested serializer")
 @pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"])
 @pytest.mark.parametrize("dtype_input_endian", [">u2", "<u2"])

tests/test_codecs/test_gzip.py

@@ -2,8 +2,12 @@
 import pytest
 
 import zarr
+from zarr.abc.codec import SupportsSyncCodec
 from zarr.abc.store import Store
 from zarr.codecs import GzipCodec
+from zarr.core.array_spec import ArrayConfig, ArraySpec
+from zarr.core.buffer import default_buffer_prototype
+from zarr.core.dtype import get_data_type_from_native_dtype
 from zarr.storage import StorePath
 
 
@@ -22,3 +26,27 @@ def test_gzip(store: Store) -> None:
 
     a[:, :] = data
     assert np.array_equal(data, a[:, :])
+
+
+def test_gzip_codec_supports_sync() -> None:
+    assert isinstance(GzipCodec(), SupportsSyncCodec)
+
+
+def test_gzip_codec_sync_roundtrip() -> None:
+    codec = GzipCodec(level=1)
+    arr = np.arange(100, dtype="float64")
+    zdtype = get_data_type_from_native_dtype(arr.dtype)
+    spec = ArraySpec(
+        shape=arr.shape,
+        dtype=zdtype,
+        fill_value=zdtype.cast_scalar(0),
+        config=ArrayConfig(order="C", write_empty_chunks=True),
+        prototype=default_buffer_prototype(),
+    )
+    buf = default_buffer_prototype().buffer.from_array_like(arr.view("B"))
+
+    encoded = codec._encode_sync(buf, spec)
+    assert encoded is not None
+    decoded = codec._decode_sync(encoded, spec)
+    result = np.frombuffer(decoded.as_numpy_array(), dtype="float64")
+    np.testing.assert_array_equal(arr, result)

tests/test_codecs/test_transpose.py

@@ -3,9 +3,13 @@
 
 import zarr
 from zarr import AsyncArray, config
+from zarr.abc.codec import SupportsSyncCodec
 from zarr.abc.store import Store
 from zarr.codecs import TransposeCodec
+from zarr.core.array_spec import ArrayConfig, ArraySpec
+from zarr.core.buffer import NDBuffer, default_buffer_prototype
 from zarr.core.common import MemoryOrder
+from zarr.core.dtype import get_data_type_from_native_dtype
 from zarr.storage import StorePath
 
 from .test_codecs import _AsyncArrayProxy
@@ -93,3 +97,27 @@ def test_transpose_invalid(
                 chunk_key_encoding={"name": "v2", "separator": "."},
                 filters=[TransposeCodec(order=order)],  # type: ignore[arg-type]
             )
+
+
+def test_transpose_codec_supports_sync() -> None:
+    assert isinstance(TransposeCodec(order=(0, 1)), SupportsSyncCodec)
+
+
+def test_transpose_codec_sync_roundtrip() -> None:
+    codec = TransposeCodec(order=(1, 0))
+    arr = np.arange(12, dtype="float64").reshape(3, 4)
+    zdtype = get_data_type_from_native_dtype(arr.dtype)
+    spec = ArraySpec(
+        shape=arr.shape,
+        dtype=zdtype,
+        fill_value=zdtype.cast_scalar(0),
+        config=ArrayConfig(order="C", write_empty_chunks=True),
+        prototype=default_buffer_prototype(),
+    )
+    nd_buf: NDBuffer = default_buffer_prototype().nd_buffer.from_numpy_array(arr)
+
+    encoded = codec._encode_sync(nd_buf, spec)
+    assert encoded is not None
+    resolved_spec = codec.resolve_metadata(spec)
+    decoded = codec._decode_sync(encoded, resolved_spec)
+    np.testing.assert_array_equal(arr, decoded.as_numpy_array())

tests/test_codecs/test_vlen.py

@@ -5,9 +5,10 @@
 
 import zarr
 from zarr import Array
-from zarr.abc.codec import Codec
+from zarr.abc.codec import Codec, SupportsSyncCodec
 from zarr.abc.store import Store
 from zarr.codecs import ZstdCodec
+from zarr.codecs.vlen_utf8 import VLenBytesCodec, VLenUTF8Codec
 from zarr.core.dtype import get_data_type_from_native_dtype
 from zarr.core.dtype.npy.string import _NUMPY_SUPPORTS_VLEN_STRING
 from zarr.core.metadata.v3 import ArrayV3Metadata
@@ -62,3 +63,11 @@ def test_vlen_string(
     assert np.array_equal(data, b[:, :])
     assert b.metadata.data_type == get_data_type_from_native_dtype(data.dtype)
     assert a.dtype == data.dtype
+
+
+def test_vlen_utf8_codec_supports_sync() -> None:
+    assert isinstance(VLenUTF8Codec(), SupportsSyncCodec)
+
+
+def test_vlen_bytes_codec_supports_sync() -> None:
+    assert isinstance(VLenBytesCodec(), SupportsSyncCodec)