use chunktransform

Author: d-v-b

src/zarr/abc/codec.py

@@ -83,9 +83,11 @@ def _encode_sync(
 class SupportsChunkCodec(Protocol):
     """Protocol for objects that can decode/encode whole chunks synchronously.
 
-    [`CodecChain`][zarr.core.codec_pipeline.CodecChain] satisfies this protocol.
+    [`ChunkTransform`][zarr.core.codec_pipeline.ChunkTransform] satisfies this protocol.
     """
 
+    array_spec: ArraySpec
+
     def decode_chunk(self, chunk_bytes: Buffer) -> NDBuffer: ...
 
     def encode_chunk(self, chunk_array: NDBuffer) -> Buffer | None: ...
@@ -316,7 +318,6 @@ def serialize(
     def prepare_read_sync(
         self,
         byte_getter: Any,
-        chunk_spec: ArraySpec,
         chunk_selection: SelectorTuple,
         codec_chain: SupportsChunkCodec,
     ) -> NDBuffer | None:
@@ -328,21 +329,19 @@ def prepare_read_sync(
         byte_getter : Any
             An object supporting ``get_sync`` (e.g.
             [`StorePath`][zarr.storage._common.StorePath]).
-        chunk_spec : ArraySpec
-            The [`ArraySpec`][zarr.core.array_spec.ArraySpec] for the chunk.
         chunk_selection : SelectorTuple
             Selection within the decoded chunk array.
         codec_chain : SupportsChunkCodec
             The [`SupportsChunkCodec`][zarr.abc.codec.SupportsChunkCodec] used to
-            decode the chunk.
+            decode the chunk. Must carry an ``array_spec`` attribute.
 
         Returns
         -------
         NDBuffer or None
             The decoded chunk data at *chunk_selection*, or ``None`` if the
             chunk does not exist in the store.
         """
-        raw = byte_getter.get_sync(prototype=chunk_spec.prototype)
+        raw = byte_getter.get_sync(prototype=codec_chain.array_spec.prototype)
         if raw is None:
             return None
         chunk_array = codec_chain.decode_chunk(raw)
@@ -351,7 +350,7 @@ def prepare_read_sync(
     def prepare_write_sync(
         self,
         byte_setter: Any,
-        chunk_spec: ArraySpec,
+        codec_chain: SupportsChunkCodec,
         chunk_selection: SelectorTuple,
         out_selection: SelectorTuple,
         replace: bool,
@@ -367,8 +366,9 @@ def prepare_write_sync(
         byte_setter : Any
             An object supporting ``get_sync`` and ``set_sync`` (e.g.
             [`StorePath`][zarr.storage._common.StorePath]).
-        chunk_spec : ArraySpec
-            The [`ArraySpec`][zarr.core.array_spec.ArraySpec] for the chunk.
+        codec_chain : SupportsChunkCodec
+            The [`SupportsChunkCodec`][zarr.abc.codec.SupportsChunkCodec]
+            carrying the ``array_spec`` for the chunk.
         chunk_selection : SelectorTuple
             Selection within the chunk being written.
         out_selection : SelectorTuple
@@ -384,6 +384,7 @@ def prepare_write_sync(
             A [`PreparedWrite`][zarr.abc.codec.PreparedWrite] carrying the
             deserialized chunk data and selection metadata.
         """
+        chunk_spec = codec_chain.array_spec
         existing: Buffer | None = None
         if not replace:
             existing = byte_setter.get_sync(prototype=chunk_spec.prototype)
@@ -403,7 +404,7 @@ def prepare_write_sync(
     def finalize_write_sync(
         self,
         prepared: PreparedWrite,
-        chunk_spec: ArraySpec,
+        codec_chain: SupportsChunkCodec,
         byte_setter: Any,
     ) -> None:
         """Serialize the prepared chunk data and write it to the store.
@@ -416,13 +417,14 @@ def finalize_write_sync(
         prepared : PreparedWrite
             The [`PreparedWrite`][zarr.abc.codec.PreparedWrite] returned by
             [`prepare_write_sync`][zarr.abc.codec.ArrayBytesCodec.prepare_write_sync].
-        chunk_spec : ArraySpec
-            The [`ArraySpec`][zarr.core.array_spec.ArraySpec] for the chunk.
+        codec_chain : SupportsChunkCodec
+            The [`SupportsChunkCodec`][zarr.abc.codec.SupportsChunkCodec]
+            carrying the ``array_spec`` for the chunk.
         byte_setter : Any
             An object supporting ``set_sync`` and ``delete_sync`` (e.g.
             [`StorePath`][zarr.storage._common.StorePath]).
         """
-        blob = self.serialize(prepared.chunk_dict, chunk_spec)
+        blob = self.serialize(prepared.chunk_dict, codec_chain.array_spec)
         if blob is None:
             byte_setter.delete_sync()
         else:
@@ -435,7 +437,6 @@ def finalize_write_sync(
     async def prepare_read(
         self,
         byte_getter: Any,
-        chunk_spec: ArraySpec,
         chunk_selection: SelectorTuple,
         codec_chain: SupportsChunkCodec,
     ) -> NDBuffer | None:
@@ -447,21 +448,19 @@ async def prepare_read(
         byte_getter : Any
             An object supporting ``get`` (e.g.
             [`StorePath`][zarr.storage._common.StorePath]).
-        chunk_spec : ArraySpec
-            The [`ArraySpec`][zarr.core.array_spec.ArraySpec] for the chunk.
         chunk_selection : SelectorTuple
             Selection within the decoded chunk array.
         codec_chain : SupportsChunkCodec
             The [`SupportsChunkCodec`][zarr.abc.codec.SupportsChunkCodec] used to
-            decode the chunk.
+            decode the chunk. Must carry an ``array_spec`` attribute.
 
         Returns
         -------
         NDBuffer or None
             The decoded chunk data at *chunk_selection*, or ``None`` if the
             chunk does not exist in the store.
         """
-        raw = await byte_getter.get(prototype=chunk_spec.prototype)
+        raw = await byte_getter.get(prototype=codec_chain.array_spec.prototype)
         if raw is None:
             return None
         chunk_array = codec_chain.decode_chunk(raw)
@@ -470,7 +469,7 @@ async def prepare_read(
     async def prepare_write(
         self,
         byte_setter: Any,
-        chunk_spec: ArraySpec,
+        codec_chain: SupportsChunkCodec,
         chunk_selection: SelectorTuple,
         out_selection: SelectorTuple,
         replace: bool,
@@ -483,8 +482,9 @@ async def prepare_write(
         byte_setter : Any
             An object supporting ``get`` and ``set`` (e.g.
             [`StorePath`][zarr.storage._common.StorePath]).
-        chunk_spec : ArraySpec
-            The [`ArraySpec`][zarr.core.array_spec.ArraySpec] for the chunk.
+        codec_chain : SupportsChunkCodec
+            The [`SupportsChunkCodec`][zarr.abc.codec.SupportsChunkCodec]
+            carrying the ``array_spec`` for the chunk.
         chunk_selection : SelectorTuple
             Selection within the chunk being written.
         out_selection : SelectorTuple
@@ -500,6 +500,7 @@ async def prepare_write(
             A [`PreparedWrite`][zarr.abc.codec.PreparedWrite] carrying the
             deserialized chunk data and selection metadata.
         """
+        chunk_spec = codec_chain.array_spec
         existing: Buffer | None = None
         if not replace:
             existing = await byte_setter.get(prototype=chunk_spec.prototype)
@@ -519,7 +520,7 @@ async def prepare_write(
     async def finalize_write(
         self,
         prepared: PreparedWrite,
-        chunk_spec: ArraySpec,
+        codec_chain: SupportsChunkCodec,
         byte_setter: Any,
     ) -> None:
         """Async variant of
@@ -530,13 +531,14 @@ async def finalize_write(
         prepared : PreparedWrite
             The [`PreparedWrite`][zarr.abc.codec.PreparedWrite] returned by
             [`prepare_write`][zarr.abc.codec.ArrayBytesCodec.prepare_write].
-        chunk_spec : ArraySpec
-            The [`ArraySpec`][zarr.core.array_spec.ArraySpec] for the chunk.
+        codec_chain : SupportsChunkCodec
+            The [`SupportsChunkCodec`][zarr.abc.codec.SupportsChunkCodec]
+            carrying the ``array_spec`` for the chunk.
         byte_setter : Any
             An object supporting ``set`` and ``delete`` (e.g.
             [`StorePath`][zarr.storage._common.StorePath]).
         """
-        blob = self.serialize(prepared.chunk_dict, chunk_spec)
+        blob = self.serialize(prepared.chunk_dict, codec_chain.array_spec)
         if blob is None:
             await byte_setter.delete()
         else:

src/zarr/core/codec_pipeline.py

@@ -69,49 +69,55 @@ def fill_value_or_default(chunk_spec: ArraySpec) -> Any:
         return fill_value
 
 
-@dataclass(frozen=True, slots=True)
-class CodecChain:
-    """Codec chain with pre-resolved metadata specs.
+@dataclass(slots=True, kw_only=True)
+class ChunkTransform:
+    """A stored chunk, modeled as a layered array.
 
-    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.
+    Each layer corresponds to one ArrayArrayCodec and the ArraySpec
+    at its input boundary.  ``layers[0]`` is the outermost (user-visible)
+    transform; after the last layer comes the ArrayBytesCodec.
+
+    The chunk's ``shape`` and ``dtype`` reflect the representation
+    **after** all ArrayArrayCodec layers have been applied — i.e. the
+    spec that feeds the ArrayBytesCodec.
     """
 
     codecs: tuple[Codec, ...]
-    chunk_spec: ArraySpec
+    array_spec: ArraySpec
 
-    _aa_codecs: tuple[ArrayArrayCodec, ...] = field(init=False, repr=False, compare=False)
-    _aa_specs: tuple[ArraySpec, ...] = field(init=False, repr=False, compare=False)
+    # Each element is (ArrayArrayCodec, input_spec_for_that_codec).
+    layers: 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)
-    _bb_spec: ArraySpec = field(init=False, repr=False, compare=False)
     _all_sync: bool = field(init=False, repr=False, compare=False)
 
     def __post_init__(self) -> None:
         aa, ab, bb = codecs_from_list(list(self.codecs))
 
-        aa_specs: list[ArraySpec] = []
-        spec = self.chunk_spec
+        layers: tuple[tuple[ArrayArrayCodec, ArraySpec], ...] = ()
+        spec = self.array_spec
         for aa_codec in aa:
-            aa_specs.append(spec)
+            layers = (*layers, (aa_codec, spec))
             spec = aa_codec.resolve_metadata(spec)
 
-        object.__setattr__(self, "_aa_codecs", aa)
-        object.__setattr__(self, "_aa_specs", tuple(aa_specs))
-        object.__setattr__(self, "_ab_codec", ab)
-        object.__setattr__(self, "_ab_spec", spec)
+        self.layers = layers
+        self._ab_codec = ab
+        self._ab_spec = spec
+        self._bb_codecs = bb
+        self._all_sync = all(isinstance(c, SupportsSyncCodec) for c in self.codecs)
 
-        spec = ab.resolve_metadata(spec)
-        object.__setattr__(self, "_bb_codecs", bb)
-        object.__setattr__(self, "_bb_spec", spec)
+    @property
+    def shape(self) -> tuple[int, ...]:
+        """Shape after all ArrayArrayCodec layers (input to the ArrayBytesCodec)."""
+        return self._ab_spec.shape
 
-        object.__setattr__(
-            self,
-            "_all_sync",
-            all(isinstance(c, SupportsSyncCodec) for c in self.codecs),
-        )
+    @property
+    def dtype(self) -> ZDType[TBaseDType, TBaseScalar]:
+        """Dtype after all ArrayArrayCodec layers (input to the ArrayBytesCodec)."""
+        return self._ab_spec.dtype
 
     @property
     def all_sync(self) -> bool:
@@ -127,11 +133,11 @@ def decode_chunk(
         """
         bb_out: Any = chunk_bytes
         for bb_codec in reversed(self._bb_codecs):
-            bb_out = cast("SupportsSyncCodec", bb_codec)._decode_sync(bb_out, self._bb_spec)
+            bb_out = cast("SupportsSyncCodec", bb_codec)._decode_sync(bb_out, self._ab_spec)
 
         ab_out: Any = cast("SupportsSyncCodec", self._ab_codec)._decode_sync(bb_out, self._ab_spec)
 
-        for aa_codec, spec in zip(reversed(self._aa_codecs), reversed(self._aa_specs), strict=True):
+        for aa_codec, spec in reversed(self.layers):
             ab_out = cast("SupportsSyncCodec", aa_codec)._decode_sync(ab_out, spec)
 
         return ab_out  # type: ignore[no-any-return]
@@ -146,7 +152,7 @@ def encode_chunk(
         """
         aa_out: Any = chunk_array
 
-        for aa_codec, spec in zip(self._aa_codecs, self._aa_specs, strict=True):
+        for aa_codec, spec in self.layers:
             if aa_out is None:
                 return None
             aa_out = cast("SupportsSyncCodec", aa_codec)._encode_sync(aa_out, spec)
@@ -158,7 +164,7 @@ def encode_chunk(
         for bb_codec in self._bb_codecs:
             if bb_out is None:
                 return None
-            bb_out = cast("SupportsSyncCodec", bb_codec)._encode_sync(bb_out, self._bb_spec)
+            bb_out = cast("SupportsSyncCodec", bb_codec)._encode_sync(bb_out, self._ab_spec)
 
         return bb_out  # type: ignore[no-any-return]
 
@@ -369,7 +375,7 @@ async def read_batch(
                     out[out_selection] = fill_value_or_default(chunk_spec)
         else:
             chunk_bytes_batch = await concurrent_map(
-                [(byte_getter, array_spec.prototype) for byte_getter, array_spec, *_ in batch_info],
+                [(byte_getter, chunk_spec.prototype) for byte_getter, chunk_spec, *_ in batch_info],
                 lambda byte_getter, prototype: byte_getter.get(prototype),
                 config.get("async.concurrency"),
             )