Merge branch 'main' into claude/zarr-metadata-api-docs-93fa91

Author: d-v-b

changes/4201.bugfix.md

@@ -0,0 +1 @@
+Fixed the opt-in `FusedCodecPipeline` for serializers that advertise the partial-decode/encode mixins with only the documented async partial methods: the partial dispatch previously asserted on the private `_decode_partial_sync`/`_encode_partial_sync` hooks (an `AssertionError`, or an `AttributeError` mid-IO under `python -O`); such codecs now take the full-chunk sync path.

changes/4202.bugfix.md

@@ -0,0 +1,10 @@
+Fixed `FusedCodecPipeline` (the opt-in synchronous pipeline) silently skipping
+array-array/bytes-bytes codecs placed outside a sharding serializer on its
+partial-decode/partial-encode fast paths. With an outer compressor (e.g.
+`compressors=[GzipCodec()]` around a `ShardingCodec` serializer), the fused
+pipeline wrote non-conforming stored bytes that `BatchedCodecPipeline` (and any
+other conforming reader) could not read, and could fail to read data that
+`BatchedCodecPipeline` had written. With an outer array-array codec (e.g.
+`TransposeCodec`), it silently returned wrong data in both directions with no
+error. Only the opt-in `FusedCodecPipeline` was affected; the default
+`BatchedCodecPipeline` was never impacted.

src/zarr/core/codec_pipeline.py

@@ -143,10 +143,10 @@ def pipeline_supports_partial_decode(
     selection non-contiguous, a BB codec can rewrite the bytes), making partial
     decode infeasible.
 
-    NOTE: the two pipelines currently pass different ``require_no_aa_bb`` values
-    (Batched: True; Fused: False). That divergence is intentional-for-now and
-    tracked separately; this function centralizes the predicate without changing
-    either pipeline's behavior.
+    Both pipelines pass `require_no_aa_bb=True`: an outer AA/BB codec (e.g. a
+    compressor wrapping a sharding serializer) must see every byte of the
+    chunk, so a partial branch that only re-decodes/re-encodes the inner
+    sharding codec would silently bypass it.
     """
     if require_no_aa_bb and (len(array_array_codecs) + len(bytes_bytes_codecs)) != 0:
         return False
@@ -162,8 +162,7 @@ def pipeline_supports_partial_encode(
 ) -> bool:
     """Whether a codec pipeline can encode a partial selection without a full rewrite.
 
-    Mirror of ``pipeline_supports_partial_decode`` for encoding. See its note re:
-    the per-pipeline ``require_no_aa_bb`` divergence.
+    Mirror of `pipeline_supports_partial_decode` for encoding.
     """
     if require_no_aa_bb and (len(array_array_codecs) + len(bytes_bytes_codecs)) != 0:
         return False
@@ -934,14 +933,11 @@ def __iter__(self) -> Iterator[Codec]:
 
     @property
     def supports_partial_decode(self) -> bool:
-        # NOTE: unlike BatchedCodecPipeline this does NOT require the AA/BB codec
-        # lists to be empty (require_no_aa_bb=False). That divergence is tracked
-        # separately; see pipeline_supports_partial_decode.
         return pipeline_supports_partial_decode(
             self.array_bytes_codec,
             array_array_codecs=self.array_array_codecs,
             bytes_bytes_codecs=self.bytes_bytes_codecs,
-            require_no_aa_bb=False,
+            require_no_aa_bb=True,
         )
 
     @property
@@ -950,7 +946,7 @@ def supports_partial_encode(self) -> bool:
             self.array_bytes_codec,
             array_array_codecs=self.array_array_codecs,
             bytes_bytes_codecs=self.bytes_bytes_codecs,
-            require_no_aa_bb=False,
+            require_no_aa_bb=True,
         )
 
     def validate(
@@ -1039,10 +1035,14 @@ def read_sync(
 
         # Partial-decode fast path: the AB codec owns IO (read only the
         # byte ranges needed for the requested selection). Same condition
-        # and dispatch as BatchedCodecPipeline.read_batch.
-        if self.supports_partial_decode:
-            codec = self.array_bytes_codec
-            assert hasattr(codec, "_decode_partial_sync")
+        # and dispatch as BatchedCodecPipeline.read_batch, plus a gate on the
+        # sync partial method: the public partial-decode contract
+        # (`ArrayBytesCodecPartialDecodeMixin`) only requires the async
+        # `_decode_partial_single`, so a codec may support partial decode
+        # without `_decode_partial_sync` — such codecs take the full-chunk
+        # path below instead.
+        codec = self.array_bytes_codec
+        if self.supports_partial_decode and hasattr(codec, "_decode_partial_sync"):
 
             def _read_one(
                 item: tuple[Any, ArraySpec, SelectorTuple, SelectorTuple, bool],
@@ -1111,10 +1111,14 @@ def write_sync(
 
         # Partial-encode path: the AB codec owns IO (read, merge, encode,
         # write).  Same condition and calling convention as
-        # BatchedCodecPipeline.write_batch.
-        if self.supports_partial_encode:
-            codec = self.array_bytes_codec
-            assert hasattr(codec, "_encode_partial_sync")
+        # BatchedCodecPipeline.write_batch, plus a gate on the sync partial
+        # method: the public partial-encode contract
+        # (`ArrayBytesCodecPartialEncodeMixin`) only requires the async
+        # `_encode_partial_single`, so a codec may support partial encode
+        # without `_encode_partial_sync` — such codecs take the full-chunk
+        # path below instead.
+        codec = self.array_bytes_codec
+        if self.supports_partial_encode and hasattr(codec, "_encode_partial_sync"):
             scalar = len(value.shape) == 0
 
             def _write_one(

tests/test_fused_pipeline.py

@@ -2,21 +2,32 @@
 
 from __future__ import annotations
 
-from typing import Any
+from dataclasses import dataclass, field, replace
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import pytest
 
 import zarr
-from zarr.abc.codec import BytesBytesCodec
+from zarr.abc.codec import (
+    ArrayBytesCodec,
+    ArrayBytesCodecPartialDecodeMixin,
+    ArrayBytesCodecPartialEncodeMixin,
+    BytesBytesCodec,
+)
 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.codec_pipeline import FusedCodecPipeline
 from zarr.core.config import config as zarr_config
+from zarr.registry import register_codec
 from zarr.storage import MemoryStore, StorePath
 
+if TYPE_CHECKING:
+    from zarr.core.array_spec import ArraySpec
+    from zarr.core.buffer import Buffer, NDBuffer
+
 
 @pytest.mark.parametrize(
     "codecs",
@@ -261,7 +272,7 @@ def test_chunk_transform_uses_runtime_prototype() -> None:
     """
     from zarr.abc.codec import BytesBytesCodec
     from zarr.core.array_spec import ArrayConfig, ArraySpec
-    from zarr.core.buffer import Buffer, BufferPrototype, default_buffer_prototype
+    from zarr.core.buffer import BufferPrototype, default_buffer_prototype
     from zarr.core.chunk_utils import ChunkTransform
     from zarr.core.dtype import get_data_type_from_native_dtype
 
@@ -831,3 +842,127 @@ def test_async_decode_encode_passes_through_none_chunks() -> None:
     assert decoded[1] is None
     assert decoded[0] is not None
     np.testing.assert_array_equal(decoded[0].as_numpy_array(), data)
+
+
+# ---------------------------------------------------------------------------
+# Graceful fallback for partial-mixin codecs without private sync-partial hooks
+#
+# The public partial-decode/encode contract (`ArrayBytesCodecPartialDecodeMixin`
+# / `ArrayBytesCodecPartialEncodeMixin`) only requires the async
+# `_decode_partial_single` / `_encode_partial_single`. The fused pipeline must
+# route such codecs through its full-chunk sync path instead of asserting on
+# the private `_decode_partial_sync` / `_encode_partial_sync` hooks. The double
+# below is a minimal conforming implementer of that contract; it guards the
+# public extension API, so it must not grow the private sync-partial methods.
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class PartialMixinCodec(
+    ArrayBytesCodec, ArrayBytesCodecPartialDecodeMixin, ArrayBytesCodecPartialEncodeMixin
+):
+    """Serializer with sync whole-chunk methods plus ONLY async partial methods.
+
+    This is the pre-fused public contract for partial-capable codecs: the
+    mixins' `_decode_partial_single` / `_encode_partial_single`. It must not
+    implement `_decode_partial_sync` / `_encode_partial_sync`.
+    """
+
+    inner: BytesCodec = field(default_factory=BytesCodec)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> PartialMixinCodec:
+        return cls()
+
+    def to_dict(self) -> dict[str, Any]:
+        return {"name": "test-partial-mixin"}
+
+    def evolve_from_array_spec(self, array_spec: ArraySpec) -> PartialMixinCodec:
+        return replace(self, inner=self.inner.evolve_from_array_spec(array_spec))
+
+    def compute_encoded_size(self, input_byte_length: int, chunk_spec: ArraySpec) -> int:
+        return self.inner.compute_encoded_size(input_byte_length, chunk_spec)
+
+    def _decode_sync(self, chunk_bytes: Buffer, chunk_spec: ArraySpec) -> NDBuffer:
+        return self.inner._decode_sync(chunk_bytes, chunk_spec)
+
+    def _encode_sync(self, chunk_array: NDBuffer, chunk_spec: ArraySpec) -> Buffer | None:
+        return self.inner._encode_sync(chunk_array, chunk_spec)
+
+    async def _decode_single(self, chunk_bytes: Buffer, chunk_spec: ArraySpec) -> NDBuffer:
+        return self._decode_sync(chunk_bytes, chunk_spec)
+
+    async def _encode_single(self, chunk_array: NDBuffer, chunk_spec: ArraySpec) -> Buffer | None:
+        return self._encode_sync(chunk_array, chunk_spec)
+
+    async def _decode_partial_single(
+        self, byte_getter: Any, selection: Any, chunk_spec: ArraySpec
+    ) -> NDBuffer | None:
+        chunk_bytes = await byte_getter.get(prototype=chunk_spec.prototype)
+        if chunk_bytes is None:
+            return None
+        return self._decode_sync(chunk_bytes, chunk_spec)[selection]
+
+    async def _encode_partial_single(
+        self, byte_setter: Any, chunk_array: NDBuffer, selection: Any, chunk_spec: ArraySpec
+    ) -> None:
+        existing = await byte_setter.get(prototype=chunk_spec.prototype)
+        if existing is None:
+            full = chunk_spec.prototype.nd_buffer.create(
+                shape=chunk_spec.shape,
+                dtype=chunk_spec.dtype.to_native_dtype(),
+                fill_value=chunk_spec.fill_value,
+            )
+        else:
+            full = self._decode_sync(existing, chunk_spec)
+        full[selection] = chunk_array
+        encoded = self._encode_sync(full, chunk_spec)
+        assert encoded is not None
+        await byte_setter.set(encoded)
+
+
+register_codec("test-partial-mixin", PartialMixinCodec)
+
+_FUSED = {"codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline"}
+_BATCHED = {"codec_pipeline.path": "zarr.core.codec_pipeline.BatchedCodecPipeline"}
+
+
+@pytest.mark.filterwarnings("ignore::zarr.errors.UnstableSpecificationWarning")
+@pytest.mark.parametrize("dtype", ["uint8", "float64"])
+def test_partial_mixin_codec_async_partial_only_round_trip(dtype: str) -> None:
+    """A serializer advertising the partial mixins with only async partial
+    methods must round-trip under the fused pipeline: full write, full read,
+    partial read, partial write, plus cross-pipeline parity with
+    BatchedCodecPipeline."""
+    data = np.arange(64, dtype=dtype).reshape(8, 8)
+
+    with zarr_config.set(_FUSED):
+        store = MemoryStore()
+        arr = zarr.create_array(
+            store,
+            shape=(8, 8),
+            chunks=(4, 4),
+            dtype=dtype,
+            serializer=PartialMixinCodec(),
+            compressors=None,
+            filters=None,
+            fill_value=0,
+        )
+
+        pipeline = arr._async_array.codec_pipeline
+        assert isinstance(pipeline, FusedCodecPipeline)
+        assert pipeline.supports_partial_decode
+        assert pipeline.supports_partial_encode
+        assert pipeline.sync_transform is not None
+
+        arr[:] = data
+        np.testing.assert_array_equal(arr[:], data)
+        np.testing.assert_array_equal(arr[1:5, 2:7], data[1:5, 2:7])
+
+        expected = data.copy()
+        expected[2:6, 1:3] = 7
+        arr[2:6, 1:3] = expected[2:6, 1:3]
+        np.testing.assert_array_equal(arr[:], expected)
+
+    with zarr_config.set(_BATCHED):
+        np.testing.assert_array_equal(zarr.open_array(store, mode="r")[:], expected)