refactor codecchain

Author: d-v-b

src/zarr/core/codec_pipeline.py

@@ -69,133 +69,103 @@ def fill_value_or_default(chunk_spec: ArraySpec) -> Any:
         return fill_value
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=True, slots=True)
 class CodecChain:
-    """Lightweight codec chain: array-array -> array-bytes -> bytes-bytes.
+    """Codec chain with pre-resolved metadata specs.
 
-    Pure compute only -- no IO methods, no threading, no batching.
+    Constructed from an iterable of codecs and a chunk ArraySpec.
+    Resolves each codec against the spec so that encode/decode can
+    run without re-resolving.  Pure compute only -- no IO, no threading,
+    no batching.
     """
 
-    array_array_codecs: tuple[ArrayArrayCodec, ...]
-    array_bytes_codec: ArrayBytesCodec
-    bytes_bytes_codecs: tuple[BytesBytesCodec, ...]
+    codecs: tuple[Codec, ...]
+    chunk_spec: ArraySpec
 
-    _all_sync: bool = field(default=False, init=False, repr=False, compare=False)
+    _aa_codecs: tuple[tuple[ArrayArrayCodec, ArraySpec], ...] = field(
+        init=False, repr=False, compare=False
+    )
+    _ab_codec: ArrayBytesCodec = field(init=False, repr=False, compare=False)
+    _ab_spec: ArraySpec = field(init=False, repr=False, compare=False)
+    _bb_codecs: tuple[BytesBytesCodec, ...] = field(init=False, repr=False, compare=False)
+    _all_sync: bool = field(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
+        aa, ab, bb = codecs_from_list(list(self.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))
+        aa_pairs: list[tuple[ArrayArrayCodec, ArraySpec]] = []
+        spec = self.chunk_spec
+        for aa_codec in aa:
+            aa_pairs.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)
+        object.__setattr__(self, "_aa_codecs", tuple(aa_pairs))
+        object.__setattr__(self, "_ab_codec", ab)
+        object.__setattr__(self, "_ab_spec", 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)
+        object.__setattr__(self, "_bb_codecs", bb)
 
-        return (aa_codecs_with_spec, ab_codec_with_spec, bb_codecs_with_spec)
+        object.__setattr__(
+            self,
+            "_all_sync",
+            all(isinstance(c, SupportsSyncCodec) for c in self.codecs),
+        )
+
+    @property
+    def all_sync(self) -> bool:
+        return self._all_sync
 
     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)
+        for bb_codec in reversed(self._bb_codecs):
+            bb_out = cast("SupportsSyncCodec", bb_codec)._decode_sync(bb_out, self.chunk_spec)
 
-        ab_codec, ab_spec = ab_pair
-        ab_out: Any = cast("SupportsSyncCodec", ab_codec)._decode_sync(bb_out, ab_spec)
+        ab_out: Any = cast("SupportsSyncCodec", self._ab_codec)._decode_sync(bb_out, self._ab_spec)
 
-        for aa_codec, spec in reversed(list(aa_chain)):
+        for aa_codec, spec in reversed(self._aa_codecs):
             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:
+        for aa_codec, spec in self._aa_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)
+        bb_out: Any = cast("SupportsSyncCodec", self._ab_codec)._encode_sync(aa_out, self._ab_spec)
 
-        for bb_codec in self.bytes_bytes_codecs:
+        for bb_codec in self._bb_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)
+            bb_out = cast("SupportsSyncCodec", bb_codec)._encode_sync(bb_out, self.chunk_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:
+        for codec in self.codecs:
             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):

tests/test_sync_codec_pipeline.py

@@ -1,21 +1,18 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
 import numpy as np
-import pytest
 
 from zarr.codecs.bytes import BytesCodec
 from zarr.codecs.gzip import GzipCodec
 from zarr.codecs.transpose import TransposeCodec
 from zarr.codecs.zstd import ZstdCodec
 from zarr.core.array_spec import ArrayConfig, ArraySpec
 from zarr.core.buffer import NDBuffer, default_buffer_prototype
+from zarr.core.codec_pipeline import CodecChain
 from zarr.core.dtype import get_data_type_from_native_dtype
 
-if TYPE_CHECKING:
-    from zarr.abc.codec import Codec
-
 
 def _make_array_spec(shape: tuple[int, ...], dtype: np.dtype[np.generic]) -> ArraySpec:
     zdtype = get_data_type_from_native_dtype(dtype)
@@ -33,124 +30,50 @@ def _make_nd_buffer(arr: np.ndarray[Any, np.dtype[Any]]) -> NDBuffer:
 
 
 class TestCodecChain:
-    def test_from_codecs_bytes_only(self) -> None:
-        from zarr.core.codec_pipeline import CodecChain
-
-        chain = CodecChain.from_codecs([BytesCodec()])
-        assert chain.array_array_codecs == ()
-        assert isinstance(chain.array_bytes_codec, BytesCodec)
-        assert chain.bytes_bytes_codecs == ()
-        assert chain._all_sync is True
-
-    def test_from_codecs_with_compression(self) -> None:
-        from zarr.core.codec_pipeline import CodecChain
-
-        chain = CodecChain.from_codecs([BytesCodec(), GzipCodec()])
-        assert isinstance(chain.array_bytes_codec, BytesCodec)
-        assert len(chain.bytes_bytes_codecs) == 1
-        assert isinstance(chain.bytes_bytes_codecs[0], GzipCodec)
-        assert chain._all_sync is True
-
-    def test_from_codecs_with_transpose(self) -> None:
-        from zarr.core.codec_pipeline import CodecChain
-
-        chain = CodecChain.from_codecs([TransposeCodec(order=(1, 0)), BytesCodec()])
-        assert len(chain.array_array_codecs) == 1
-        assert isinstance(chain.array_array_codecs[0], TransposeCodec)
-        assert isinstance(chain.array_bytes_codec, BytesCodec)
-        assert chain._all_sync is True
+    def test_all_sync(self) -> None:
+        spec = _make_array_spec((100,), np.dtype("float64"))
+        chain = CodecChain((BytesCodec(),), spec)
+        assert chain.all_sync is True
 
-    def test_from_codecs_full_chain(self) -> None:
-        from zarr.core.codec_pipeline import CodecChain
+    def test_all_sync_with_compression(self) -> None:
+        spec = _make_array_spec((100,), np.dtype("float64"))
+        chain = CodecChain((BytesCodec(), GzipCodec()), spec)
+        assert chain.all_sync is True
 
-        chain = CodecChain.from_codecs([TransposeCodec(order=(1, 0)), BytesCodec(), ZstdCodec()])
-        assert len(chain.array_array_codecs) == 1
-        assert isinstance(chain.array_bytes_codec, BytesCodec)
-        assert len(chain.bytes_bytes_codecs) == 1
-        assert chain._all_sync is True
-
-    def test_iter(self) -> None:
-        from zarr.core.codec_pipeline import CodecChain
-
-        codecs: list[Codec] = [TransposeCodec(order=(1, 0)), BytesCodec(), GzipCodec()]
-        chain = CodecChain.from_codecs(codecs)
-        assert list(chain) == codecs
-
-    def test_frozen(self) -> None:
-        from zarr.core.codec_pipeline import CodecChain
-
-        chain = CodecChain.from_codecs([BytesCodec()])
-        with pytest.raises(AttributeError):
-            chain.array_bytes_codec = BytesCodec()  # type: ignore[misc]
+    def test_all_sync_full_chain(self) -> None:
+        spec = _make_array_spec((3, 4), np.dtype("float64"))
+        chain = CodecChain((TransposeCodec(order=(1, 0)), BytesCodec(), ZstdCodec()), spec)
+        assert chain.all_sync is True
 
     def test_encode_decode_roundtrip_bytes_only(self) -> None:
-        from zarr.core.codec_pipeline import CodecChain
-
-        chain = CodecChain.from_codecs([BytesCodec()])
         arr = np.arange(100, dtype="float64")
         spec = _make_array_spec(arr.shape, arr.dtype)
-        chain_evolved = CodecChain.from_codecs([c.evolve_from_array_spec(spec) for c in chain])
+        chain = CodecChain((BytesCodec(),), spec)
         nd_buf = _make_nd_buffer(arr)
 
-        encoded = chain_evolved.encode_chunk(nd_buf, spec)
+        encoded = chain.encode_chunk(nd_buf)
         assert encoded is not None
-        decoded = chain_evolved.decode_chunk(encoded, spec)
-        assert decoded is not None
+        decoded = chain.decode_chunk(encoded)
         np.testing.assert_array_equal(arr, decoded.as_numpy_array())
 
     def test_encode_decode_roundtrip_with_compression(self) -> None:
-        from zarr.core.codec_pipeline import CodecChain
-
-        chain = CodecChain.from_codecs([BytesCodec(), GzipCodec(level=1)])
         arr = np.arange(100, dtype="float64")
         spec = _make_array_spec(arr.shape, arr.dtype)
-        chain_evolved = CodecChain.from_codecs([c.evolve_from_array_spec(spec) for c in chain])
+        chain = CodecChain((BytesCodec(), GzipCodec(level=1)), spec)
         nd_buf = _make_nd_buffer(arr)
 
-        encoded = chain_evolved.encode_chunk(nd_buf, spec)
+        encoded = chain.encode_chunk(nd_buf)
         assert encoded is not None
-        decoded = chain_evolved.decode_chunk(encoded, spec)
-        assert decoded is not None
+        decoded = chain.decode_chunk(encoded)
         np.testing.assert_array_equal(arr, decoded.as_numpy_array())
 
     def test_encode_decode_roundtrip_with_transpose(self) -> None:
-        from zarr.core.codec_pipeline import CodecChain
-
-        chain = CodecChain.from_codecs(
-            [TransposeCodec(order=(1, 0)), BytesCodec(), ZstdCodec(level=1)]
-        )
         arr = np.arange(12, dtype="float64").reshape(3, 4)
         spec = _make_array_spec(arr.shape, arr.dtype)
-        chain_evolved = CodecChain.from_codecs([c.evolve_from_array_spec(spec) for c in chain])
+        chain = CodecChain((TransposeCodec(order=(1, 0)), BytesCodec(), ZstdCodec(level=1)), spec)
         nd_buf = _make_nd_buffer(arr)
 
-        encoded = chain_evolved.encode_chunk(nd_buf, spec)
+        encoded = chain.encode_chunk(nd_buf)
         assert encoded is not None
-        decoded = chain_evolved.decode_chunk(encoded, spec)
-        assert decoded is not None
+        decoded = chain.decode_chunk(encoded)
         np.testing.assert_array_equal(arr, decoded.as_numpy_array())
-
-    def test_resolve_metadata_chain(self) -> None:
-        from zarr.core.codec_pipeline import CodecChain
-
-        chain = CodecChain.from_codecs([TransposeCodec(order=(1, 0)), BytesCodec(), GzipCodec()])
-        arr = np.zeros((3, 4), dtype="float64")
-        spec = _make_array_spec(arr.shape, arr.dtype)
-        chain_evolved = CodecChain.from_codecs([c.evolve_from_array_spec(spec) for c in chain])
-
-        aa_chain, ab_pair, bb_chain = chain_evolved.resolve_metadata_chain(spec)
-        assert len(aa_chain) == 1
-        assert aa_chain[0][1].shape == (3, 4)  # spec before transpose
-        _ab_codec, ab_spec = ab_pair
-        assert ab_spec.shape == (4, 3)  # spec after transpose
-        assert len(bb_chain) == 1
-
-    def test_resolve_metadata(self) -> None:
-        from zarr.core.codec_pipeline import CodecChain
-
-        chain = CodecChain.from_codecs([TransposeCodec(order=(1, 0)), BytesCodec()])
-        spec = _make_array_spec((3, 4), np.dtype("float64"))
-        chain_evolved = CodecChain.from_codecs([c.evolve_from_array_spec(spec) for c in chain])
-        resolved = chain_evolved.resolve_metadata(spec)
-        # After transpose (1,0) + bytes, shape should reflect the transpose
-        assert resolved.shape == (4, 3)