From 3afccd2f35609c7d6da658963eb18aa4d02035d2 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 29 Jul 2026 13:21:20 +0200 Subject: [PATCH] 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])