Skip to content
Open
1 change: 1 addition & 0 deletions docs/release-notes.md
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down
14 changes: 14 additions & 0 deletions src/zarr/abc/codec.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down
26 changes: 24 additions & 2 deletions src/zarr/codecs/sharding.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
ArrayBytesCodecPartialEncodeMixin,
Codec,
CodecPipeline,
SupportsSyncCodec,
_codec_supports_sync,
)
from zarr.abc.store import (
ByteGetter,
Expand Down Expand Up @@ -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, ...]
Expand Down
10 changes: 8 additions & 2 deletions src/zarr/core/chunk_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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(
Expand Down
88 changes: 88 additions & 0 deletions tests/test_fused_pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
Loading