From 3afccd2f35609c7d6da658963eb18aa4d02035d2 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 29 Jul 2026 13:21:20 +0200 Subject: [PATCH 1/2] fix: FusedCodecPipeline must apply outer AA/BB codecs on partial paths FusedCodecPipeline.supports_partial_decode/supports_partial_encode passed require_no_aa_bb=False, unlike BatchedCodecPipeline (True). With an outer array-array or bytes-bytes codec around a sharding serializer (e.g. compressors=[GzipCodec()], or filters=[TransposeCodec()]), the fused pipeline's partial read/write branches called ShardingCodec's partial sync methods directly on the raw stored value, skipping those outer codecs entirely. That wrote non-conforming bytes for an outer BB codec (unreadable by BatchedCodecPipeline or any conforming reader) and silently produced wrong data for an outer AA codec. Pass require_no_aa_bb=True in both fused properties so these chains fall through to the full-chunk fused path instead, matching batched behavior. Adds cross-pipeline parity coverage (full and partial read/write) for sharding with an outer compressor and with an outer transpose filter, and removes the "known limitation" exclusion that previously kept the sharding+compressor case out of the nested-sharding parity matrix. Assisted-by: ClaudeCode:claude-sonnet-5 --- changes/4202.bugfix.md | 10 +++ src/zarr/core/codec_pipeline.py | 18 ++--- tests/test_pipeline_parity.py | 133 ++++++++++++++++++++++++++++---- 3 files changed, 135 insertions(+), 26 deletions(-) create mode 100644 changes/4202.bugfix.md diff --git a/changes/4202.bugfix.md b/changes/4202.bugfix.md new file mode 100644 index 0000000000..6130fc5b33 --- /dev/null +++ b/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. diff --git a/src/zarr/core/codec_pipeline.py b/src/zarr/core/codec_pipeline.py index 4b8831bc7b..12d762d260 100644 --- a/src/zarr/core/codec_pipeline.py +++ b/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( diff --git a/tests/test_pipeline_parity.py b/tests/test_pipeline_parity.py index 717f0f48f1..94d95c4c24 100644 --- a/tests/test_pipeline_parity.py +++ b/tests/test_pipeline_parity.py @@ -33,6 +33,8 @@ from __future__ import annotations +import warnings +from contextlib import contextmanager from typing import TYPE_CHECKING, Any import numpy as np @@ -48,7 +50,9 @@ ShardingCodec, SubchunkWriteOrder, ) +from zarr.codecs.transpose import TransposeCodec from zarr.core.config import config as zarr_config +from zarr.errors import ZarrUserWarning from zarr.storage import MemoryStore if TYPE_CHECKING: @@ -107,11 +111,15 @@ def _store_snapshot(store: MemoryStore) -> dict[str, bytes]: ("2d-unsharded", {"shape": (20, 20), "chunks": (5, 5), "shards": None}), ("2d-sharded", {"shape": (20, 20), "chunks": (5, 5), "shards": (10, 10)}), # Nested sharding: outer chunk (10,10) sharded into inner chunks (5,5). - # Restricted to bytes-only codec because combining an outer ShardingCodec - # with a compressor (gzip) triggers a ZarrUserWarning and results in a - # checksum mismatch inside the inner shard index — a known limitation, not - # a pipeline-parity bug. The bytes-only path still exercises the full - # two-level shard encoding/decoding in both pipelines. + # Restricted to the codec configs that don't set their own `serializer` + # (bytes-only, gzip): this layout supplies an explicit nested-ShardingCodec + # `serializer`, and a codec config that also sets `serializer` (e.g. + # bytes-big-endian) would silently clobber it via dict merge, dropping + # sharding from the test entirely rather than exercising it. The gzip + # config applies as an outer bytes-bytes codec around the outer + # ShardingCodec -- this is the regression coverage for the fused pipeline + # applying outer AA/BB codecs around sharding (see + # `pipeline_supports_partial_decode`/`pipeline_supports_partial_encode`). ( "2d-nested-sharded", { @@ -122,9 +130,7 @@ def _store_snapshot(store: MemoryStore) -> dict[str, bytes]: chunk_shape=(10, 10), codecs=[ShardingCodec(chunk_shape=(5, 5))], ), - # Only run with the bytes-only codec config; gzip is incompatible - # with nested sharding (see comment above). - "_codec_ids": {"bytes-only"}, + "_codec_ids": {"bytes-only", "gzip"}, }, ), ] @@ -226,6 +232,23 @@ def _matrix() -> Iterator[Any]: # --------------------------------------------------------------------------- +@contextmanager +def _ignore_sharding_combo_warning() -> Iterator[None]: + """Suppress the "combining sharding_indexed disables partial reads" warning. + + Only the nested-sharded-plus-outer-codec matrix cell emits this; scoping the + ignore filter to just its message/category (rather than blanket-disabling + warnings) keeps every other warning in the run promoted to an error as usual. + """ + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message=r"Combining a `sharding_indexed` codec.*", + category=ZarrUserWarning, + ) + yield + + def _write_under_pipeline( pipeline_path: str, codec_kwargs: CodecConfig, @@ -244,12 +267,13 @@ def _write_under_pipeline( create_kwargs = {"dtype": "float64", **array_layout, **codec_kwargs} store = MemoryStore() with zarr_config.set({"codec_pipeline.path": pipeline_path}): - arr = zarr.create_array( - store=store, - fill_value=0, - config={"write_empty_chunks": write_empty_chunks}, - **create_kwargs, - ) + with _ignore_sharding_combo_warning(): + arr = zarr.create_array( + store=store, + fill_value=0, + config={"write_empty_chunks": write_empty_chunks}, + **create_kwargs, + ) for sel, val in sequence: arr[sel] = val contents = arr[...] @@ -259,7 +283,8 @@ def _write_under_pipeline( def _read_under_pipeline(pipeline_path: str, store: MemoryStore) -> Any: """Re-open an existing store under the chosen pipeline and read it whole.""" with zarr_config.set({"codec_pipeline.path": pipeline_path}): - arr = zarr.open_array(store=store, mode="r") + with _ignore_sharding_combo_warning(): + arr = zarr.open_array(store=store, mode="r") return arr[...] @@ -418,3 +443,81 @@ def run(pipeline_path: str) -> tuple[dict[str, bytes], Any]: f"(index_location={index_location!r}) — byte-range write fast path likely assumed " f"the wrong physical chunk order" ) + + +# --------------------------------------------------------------------------- +# Outer array-array / bytes-bytes codecs around a sharding serializer +# --------------------------------------------------------------------------- +# +# Regression coverage for FusedCodecPipeline.supports_partial_decode/encode: +# it used to allow AA/BB codecs outside the sharding codec, so its partial +# branches called ShardingCodec._decode_partial_sync/_encode_partial_sync +# directly on the raw stored value, skipping any outer filter/compressor. +# That corrupted on-disk bytes for an outer bytes-bytes codec (unreadable by +# the other pipeline) and silently produced wrong data for an outer +# array-array codec. Both configs below force the partial branches: a +# region write and a region read are included alongside the full ones. + +_OUTER_AA_BB_CONFIGS: list[tuple[str, CodecConfig]] = [ + ( + "outer-gzip-around-sharding", + { + "serializer": ShardingCodec(chunk_shape=(2, 2)), + "compressors": [GzipCodec(level=1)], + }, + ), + ( + "outer-transpose-around-sharding", + { + "filters": [TransposeCodec(order=(1, 0))], + "serializer": ShardingCodec(chunk_shape=(2, 2)), + "compressors": None, + }, + ), +] + + +@pytest.mark.parametrize(("config_id", "codec_kwargs"), _OUTER_AA_BB_CONFIGS) +@pytest.mark.parametrize( + ("writer", "reader"), + [(_BATCHED, _FUSED), (_FUSED, _BATCHED)], + ids=["batched-write-fused-read", "fused-write-batched-read"], +) +def test_pipeline_parity_outer_aa_bb_codecs( + config_id: str, + codec_kwargs: CodecConfig, + writer: str, + reader: str, +) -> None: + """Data written under one pipeline with outer AA/BB codecs must read back + correctly under the other, including through a partial write and a + partial read. + """ + shape = (8, 8) + data = (np.arange(int(np.prod(shape))).reshape(shape) + 1).astype("uint16") + store = MemoryStore() + + with zarr_config.set({"codec_pipeline.path": writer}): + with _ignore_sharding_combo_warning(): + arr = zarr.create_array( + store=store, + shape=shape, + chunks=(4, 4), + dtype=data.dtype, + fill_value=0, + **codec_kwargs, + ) + arr[...] = data + arr[2:5, 1:3] = 99 # region write -- exercises the partial-encode branch + + expected = data.copy() + expected[2:5, 1:3] = 99 + + with zarr_config.set({"codec_pipeline.path": reader}): + with _ignore_sharding_combo_warning(): + arr2 = zarr.open_array(store=store, mode="r") + full = arr2[...] + partial = arr2[1:3, 2:7] # region read -- exercises the partial-decode branch + + np.testing.assert_array_equal(full, expected) + np.testing.assert_array_equal(partial, expected[1:3, 2:7]) From 80880a0425243406e145a69b722d7411c1dc6563 Mon Sep 17 00:00:00 2001 From: Davis Bennett Date: Wed, 29 Jul 2026 17:54:21 +0200 Subject: [PATCH 2/2] fix: fused pipeline falls back for partial-mixin codecs without sync partial methods (#4201) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The partial dispatch in FusedCodecPipeline.read_sync/write_sync asserted the private _decode_partial_sync/_encode_partial_sync methods, which only ShardingCodec implements. A codec advertising the public partial mixins (ArrayBytesCodecPartialDecodeMixin/-EncodeMixin) with only the documented async partial methods died with a bare AssertionError — or, under python -O, an AttributeError mid-IO. The asserts are now capability gates: codecs without the sync partial methods take the full-chunk sync path instead. The related crash for sharded arrays with async-only inner codecs is fixed separately in zarr-developers/zarr-python#4179. Assisted-by: ClaudeCode:claude-fable-5 --- changes/4201.bugfix.md | 1 + src/zarr/core/codec_pipeline.py | 24 ++++-- tests/test_fused_pipeline.py | 141 +++++++++++++++++++++++++++++++- 3 files changed, 155 insertions(+), 11 deletions(-) create mode 100644 changes/4201.bugfix.md diff --git a/changes/4201.bugfix.md b/changes/4201.bugfix.md new file mode 100644 index 0000000000..d837a8a9e2 --- /dev/null +++ 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. diff --git a/src/zarr/core/codec_pipeline.py b/src/zarr/core/codec_pipeline.py index 4b8831bc7b..56a06b906c 100644 --- a/src/zarr/core/codec_pipeline.py +++ b/src/zarr/core/codec_pipeline.py @@ -1039,10 +1039,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 +1115,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( diff --git a/tests/test_fused_pipeline.py b/tests/test_fused_pipeline.py index 02b4026fd9..fd86936853 100644 --- a/tests/test_fused_pipeline.py +++ b/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)