diff --git a/docs/release-notes.md b/docs/release-notes.md index 3fd8a5f360..16f92ea571 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -30,6 +30,7 @@ emits), instead of falling back to the system's native byte order. ([#3417](https://github.com/zarr-developers/zarr-python/issues/3417)) - Fixed `save_array`, `Group.__setitem__`, and `load` for 0-dimensional arrays. ([#3469](https://github.com/zarr-developers/zarr-python/issues/3469)) - Fixed inner-codec spec evolution for sharded arrays. The sharding codec now threads the array spec through its inner codec chain when evolving codecs, so a codec that changes the dtype upstream of `BytesCodec` no longer leaves the inner chain evolved against the wrong spec (which previously failed at decode time). This runs on the default `BatchedCodecPipeline` as well. Standard inner chains (`[BytesCodec]`, `[BytesCodec, ZstdCodec]`, transpose + bytes) are byte-identical to before. Restores the behavior of #2179. ([#3885](https://github.com/zarr-developers/zarr-python/issues/3885)) +- Fixed the opt-in `FusedCodecPipeline` for sharded arrays whose inner or index codec chain contains a codec implementing only the async codec interface (no `SupportsSyncCodec`). Such arrays previously raised `TypeError: All codecs must implement SupportsSyncCodec` on both read and write; the pipeline now declines its synchronous fast path for them and falls back to the async path, matching the behavior of the default `BatchedCodecPipeline`. Fully sync-capable codec chains keep the fast path unchanged. ([#4179](https://github.com/zarr-developers/zarr-python/issues/4179)) - Make chunk normalization properly handle `-1` as a compact representation of the length of an entire axis. Reject several previously-accepted but ill-defined chunk specifications: `chunks=True` (previously silently produced size-1 chunks), diff --git a/src/zarr/abc/codec.py b/src/zarr/abc/codec.py index 61c5dc9948..34d349e6d1 100644 --- a/src/zarr/abc/codec.py +++ b/src/zarr/abc/codec.py @@ -82,6 +82,20 @@ def _decode_sync(self, chunk_data: CO, chunk_spec: ArraySpec) -> CI: ... def _encode_sync(self, chunk_data: CI, chunk_spec: ArraySpec) -> CO | None: ... +def _codec_supports_sync(codec: object) -> bool: + """Whether `codec` can actually run on a synchronous (no event loop) path. + + Structural membership in `SupportsSyncCodec` is necessary but not always + sufficient: a codec can provide `_decode_sync`/`_encode_sync` whose ability + to run depends on runtime configuration the type system cannot see. + `ShardingCodec` is the canonical case — its sync methods delegate to its + configured inner and index codec chains, so they only work when every codec + in those chains is itself sync-capable. Such codecs opt out dynamically via + a `_sync_capable` attribute/property (absent means capable). + """ + return isinstance(codec, SupportsSyncCodec) and getattr(codec, "_sync_capable", True) + + class BaseCodec[CI: CodecInput, CO: CodecOutput](Metadata): """Generic base class for codecs. diff --git a/src/zarr/codecs/sharding.py b/src/zarr/codecs/sharding.py index 2d4d63d400..2a95109e6b 100644 --- a/src/zarr/codecs/sharding.py +++ b/src/zarr/codecs/sharding.py @@ -14,7 +14,7 @@ ArrayBytesCodecPartialEncodeMixin, Codec, CodecPipeline, - SupportsSyncCodec, + _codec_supports_sync, ) from zarr.abc.store import ( ByteGetter, @@ -1406,8 +1406,30 @@ def _is_complete_shard_write( is_complete_chunk for *_, is_complete_chunk in indexed_chunks ) + @property + def _sync_capable(self) -> bool: + """Dynamic opt-out consulted by `_codec_supports_sync` / `ChunkTransform`. + + This codec structurally satisfies `SupportsSyncCodec`, but every sync + method (`_decode_sync`, `_encode_sync`, `_decode_partial_sync`, + `_encode_partial_sync`) delegates to the inner and index codec chains + through `ChunkTransform`, so it can only run synchronously when every + codec in BOTH chains is itself sync-capable. Reporting False here makes + `ChunkTransform` construction raise, which in turn makes + `FusedCodecPipeline.evolve_from_array_spec` set `sync_transform=None` — + the whole pipeline then declines the sync fast path and routes through + the async paths (partial shard decode / async fallback write), exactly + as it does for an async-only TOP-level codec or a non-sync store. + """ + return self._inner_codecs_sync_capable() and self._index_codecs_sync_capable() + + def _inner_codecs_sync_capable(self) -> bool: + # _codec_supports_sync (not bare isinstance) so a nested sharding codec + # with an async-only inner chain propagates its opt-out outward. + return all(_codec_supports_sync(c) for c in self.codecs) + def _index_codecs_sync_capable(self) -> bool: - return all(isinstance(c, SupportsSyncCodec) for c in self.index_codecs) + return all(_codec_supports_sync(c) for c in self.index_codecs) async def _decode_shard_index( self, index_bytes: Buffer, chunks_per_shard: tuple[int, ...] diff --git a/src/zarr/core/chunk_utils.py b/src/zarr/core/chunk_utils.py index ee42e60cce..d93793f853 100644 --- a/src/zarr/core/chunk_utils.py +++ b/src/zarr/core/chunk_utils.py @@ -3,7 +3,7 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, cast -from zarr.abc.codec import GetResult, SupportsSyncCodec +from zarr.abc.codec import GetResult, SupportsSyncCodec, _codec_supports_sync from zarr.core.indexing import is_scalar if TYPE_CHECKING: @@ -240,7 +240,13 @@ class ChunkTransform: def __post_init__(self) -> None: from zarr.core.codec_pipeline import codecs_from_list - non_sync = [c for c in self.codecs if not isinstance(c, SupportsSyncCodec)] + # _codec_supports_sync, not a bare isinstance check: a codec can satisfy + # the SupportsSyncCodec protocol structurally yet be unable to run + # synchronously (ShardingCodec whose inner/index chain contains an + # async-only codec). Such codecs opt out via `_sync_capable`, and the + # TypeError here is what makes FusedCodecPipeline.evolve_from_array_spec + # decline the sync fast path and fall back to the async pipeline. + non_sync = [c for c in self.codecs if not _codec_supports_sync(c)] if non_sync: names = ", ".join(type(c).__name__ for c in non_sync) raise TypeError( diff --git a/tests/test_fused_pipeline.py b/tests/test_fused_pipeline.py index 73c2c6e1c3..489ba310c7 100644 --- a/tests/test_fused_pipeline.py +++ b/tests/test_fused_pipeline.py @@ -8,6 +8,7 @@ import pytest import zarr +from zarr.abc.codec import BytesBytesCodec from zarr.codecs.bytes import BytesCodec from zarr.codecs.gzip import GzipCodec from zarr.codecs.transpose import TransposeCodec @@ -582,6 +583,93 @@ def spy_write_sync(self: Any, *args: Any, **kwargs: Any) -> Any: ) +# --------------------------------------------------------------------------- +# Async-only codecs inside a shard's inner codec chain +# --------------------------------------------------------------------------- + + +class _AsyncOnlyNoopCodec(BytesBytesCodec): # type: ignore[misc,unused-ignore] + """A no-op BB codec implementing ONLY the async codec interface. + + Deliberately does NOT satisfy `SupportsSyncCodec` (no `_decode_sync` / + `_encode_sync`), modelling a third-party codec that predates the sync + protocol. Class-level counters prove the codec actually ran. + """ + + is_fixed_size = True + encode_calls = 0 + decode_calls = 0 + + def to_dict(self) -> dict[str, Any]: + return {"name": "test-async-only-noop", "configuration": {}} + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> _AsyncOnlyNoopCodec: + return cls() + + def compute_encoded_size(self, input_byte_length: int, _spec: Any) -> int: + return input_byte_length + + async def _encode_single(self, chunk_bytes: Any, chunk_spec: Any) -> Any: + type(self).encode_calls += 1 + return chunk_bytes + + async def _decode_single(self, chunk_bytes: Any, chunk_spec: Any) -> Any: + type(self).decode_calls += 1 + return chunk_bytes + + +def test_sharded_roundtrip_with_async_only_inner_codec() -> None: + """A sharded array whose INNER codec chain contains an async-only codec + round-trips under FusedCodecPipeline (full write, partial write, full read, + partial read). + + Regression: the pipeline's top-level guard (evolve_from_array_spec -> + sync_transform=None) only inspected the top-level chain. ShardingCodec + structurally satisfies SupportsSyncCodec, so a sync transform was built and + the sync fast path dove into ShardingCodec's sync shard paths, which raised + TypeError from the inner ChunkTransform. The pipeline must instead decline + the sync fast path and fall back to the async inner pipeline, like + BatchedCodecPipeline. + """ + _AsyncOnlyNoopCodec.encode_calls = 0 + _AsyncOnlyNoopCodec.decode_calls = 0 + + with zarr_config.set({"codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline"}): + store = MemoryStore() + arr = zarr.create_array( + store=store, + shape=(16, 16), + shards=(8, 8), + chunks=(4, 4), + dtype="int32", + compressors=[_AsyncOnlyNoopCodec()], + fill_value=-1, + ) + assert isinstance(arr._async_array.codec_pipeline, FusedCodecPipeline) + + data = np.arange(256, dtype="int32").reshape(16, 16) + arr[:] = data # full write + np.testing.assert_array_equal(arr[:], data) # full read + np.testing.assert_array_equal(arr[2:11, 3:14], data[2:11, 3:14]) # partial read + + arr[5:7, 5:13] = 0 # partial write (read-merge-write of existing shards) + data[5:7, 5:13] = 0 + np.testing.assert_array_equal(arr[:], data) + + assert _AsyncOnlyNoopCodec.encode_calls > 0, "async-only inner codec never encoded" + assert _AsyncOnlyNoopCodec.decode_calls > 0, "async-only inner codec never decoded" + + # The stored bytes are valid for the default pipeline too: read them back + # under BatchedCodecPipeline (default codec_pipeline.path). Opening from + # metadata needs the codec name in the registry. + from zarr.registry import register_codec + + register_codec("test-async-only-noop", _AsyncOnlyNoopCodec) + reread = zarr.open_array(store=store, mode="r") + np.testing.assert_array_equal(reread[:], data) + + # --------------------------------------------------------------------------- # AsyncChunkTransform: the async per-chunk codec chain used on the async # fallback path. It is the async mirror of ChunkTransform, so it must produce