From 4810e88c16ef40b8a4f4d7ee2072d35e151bd097 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 17 Aug 2026 17:47:09 +0200 Subject: [PATCH 01/11] fix(metadata): validate codec chains against the threaded chunk spec `ArrayV3Metadata` validated every codec against the array-level shape and chunk grid, and threaded the *array* spec (not a chunk spec) through `resolve_metadata` during evolution. Both wrongly reject chains in which an earlier array->array codec changes a chunk's shape or rank, e.g. the zarr-extensions `reshape` codec followed by `transpose` with an order of the reshaped rank -- a combination the reshape spec explicitly endorses and that the encode path already handles correctly. Codecs are now evolved and validated in a single threaded pass (`evolve_and_validate_codecs`): each codec sees the chunk spec produced by the previous codec's `resolve_metadata`, exactly as at encode time. The array-level shape/chunk grid are passed to `Codec.validate` unchanged until a codec changes the chunk shape, after which the resolved chunk shape (and a regular grid of it) stands in for them. `ShardingCodec.validate` now validates its inner chain the same way against the inner chunk shape. Assisted-by: ClaudeCode:claude-fable-5 --- changes/+codec-chain-validation.bugfix.md | 1 + src/zarr/codecs/sharding.py | 18 +++ src/zarr/core/metadata/v3.py | 97 ++++++++--- .../test_codec_chain_validation.py | 150 ++++++++++++++++++ 4 files changed, 245 insertions(+), 21 deletions(-) create mode 100644 changes/+codec-chain-validation.bugfix.md create mode 100644 tests/test_codecs/test_codec_chain_validation.py diff --git a/changes/+codec-chain-validation.bugfix.md b/changes/+codec-chain-validation.bugfix.md new file mode 100644 index 0000000000..5415125cba --- /dev/null +++ b/changes/+codec-chain-validation.bugfix.md @@ -0,0 +1 @@ +Codec chains are now validated against the chunk spec threaded through each codec's `resolve_metadata`, the same way the codec pipeline resolves it at encode time, instead of against the array-level shape. Previously, an `array -> array` codec that changes a chunk's shape or rank (such as the zarr-extensions `reshape` codec) followed by a codec whose configuration refers to the transformed chunk (e.g. `transpose`) was wrongly rejected. The sharding codec now validates its inner codec chain in the same way. diff --git a/src/zarr/codecs/sharding.py b/src/zarr/codecs/sharding.py index bd0760f7e3..058cf8dc3f 100644 --- a/src/zarr/codecs/sharding.py +++ b/src/zarr/codecs/sharding.py @@ -70,6 +70,7 @@ ChunkGridMetadata, RectilinearChunkGridMetadata, RegularChunkGridMetadata, + evolve_and_validate_codecs, parse_codecs, ) from zarr.registry import get_ndbuffer_class, get_pipeline_class @@ -613,6 +614,23 @@ def validate( f"Chunk edge length {edge} in dimension {i} is not " f"divisible by the shard's inner chunk size {inner}." ) + # The inner codecs see chunks of `self.chunk_shape`; validate them + # against that, threading the chunk spec through the chain exactly as + # the top-level metadata does (an inner reshape may change the rank + # seen by a following transpose). + evolve_and_validate_codecs( + self.codecs, + shape=self.chunk_shape, + chunk_grid=RegularChunkGridMetadata(chunk_shape=self.chunk_shape), + chunk_spec=ArraySpec( + shape=self.chunk_shape, + dtype=dtype, + fill_value=dtype.default_scalar(), + config=ArrayConfig.from_dict({}), + prototype=default_buffer_prototype(), + ), + evolve=False, + ) def _get_inner_chunk_transform(self, shard_spec: ArraySpec) -> Any: """The synchronous transform for the inner codec chain. diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index 11f3eb593d..2bb0f15486 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -125,6 +125,64 @@ def validate_codecs(codecs: tuple[Codec, ...], dtype: ZDType[TBaseDType, TBaseSc ) +def representative_chunk_shape(chunk_grid: ChunkGridMetadata) -> tuple[int, ...]: + """A single chunk shape standing in for every chunk of ``chunk_grid``. + + Regular grids have exactly one chunk shape. Rectilinear grids have many; + the largest edge along each dimension is used, which is enough for the + metadata-time uses of this value (rank checks and threading a chunk spec + through ``Codec.resolve_metadata``). + """ + if isinstance(chunk_grid, RegularChunkGridMetadata): + return chunk_grid.chunk_shape + return tuple(s if isinstance(s, int) else max(s) for s in chunk_grid.chunk_shapes) + + +def evolve_and_validate_codecs( + codecs: Iterable[Codec], + *, + shape: tuple[int, ...], + chunk_grid: ChunkGridMetadata, + chunk_spec: ArraySpec, + evolve: bool = True, +) -> tuple[Codec, ...]: + """Evolve (optionally) and validate a codec chain, threading the chunk spec. + + Each codec is evolved and validated against the chunk spec produced by the + previous codec's ``resolve_metadata`` — the same spec it will see at + encode/decode time — not against the array-level metadata. Earlier + array->array codecs may change the dtype (``cast_value``) or the shape and + even the rank of a chunk (the ``reshape`` extension codec, which the spec + explicitly allows to be followed by ``transpose``). + + ``shape`` and ``chunk_grid`` are the array-level values passed to + ``Codec.validate``. They are handed unchanged to every codec until one + changes the chunk shape; from then on the array-level values are no longer + meaningful for the remaining codecs, so they are replaced by the resolved + chunk shape and a regular grid of that shape (the only shape-related facts + that survive a per-chunk reshape). ``Codec.validate`` implementations only + inspect these for rank and divisibility, so this keeps the checks sound. + + Per-codec ``validate`` runs before ``resolve_metadata``, since the latter + may rely on invariants the former checks (e.g. ``cast_value`` rejects + complex source dtypes that would otherwise crash ``_do_cast``). + """ + out: list[Codec] = [] + spec = chunk_spec + stage_shape = shape + stage_grid = chunk_grid + for codec in codecs: + evolved = codec.evolve_from_array_spec(spec) if evolve else codec + evolved.validate(shape=stage_shape, dtype=spec.dtype, chunk_grid=stage_grid) + out.append(evolved) + next_spec = evolved.resolve_metadata(spec) + if next_spec.shape != spec.shape: + stage_shape = next_spec.shape + stage_grid = RegularChunkGridMetadata(chunk_shape=next_spec.shape) + spec = next_spec + return tuple(out) + + def parse_dimension_names(data: object) -> tuple[str | None, ...] | None: if data is None: return data @@ -519,28 +577,23 @@ def __init__( codecs_parsed_partial = parse_codecs(codecs) storage_transformers_parsed = parse_storage_transformers(storage_transformers) extra_fields_parsed = parse_extra_fields(extra_fields) - array_spec = ArraySpec( - shape=shape_parsed, + if len(shape_parsed) != chunk_grid_parsed.ndim: + raise ValueError("`chunk_grid` and `shape` need to have the same number of dimensions.") + # Codecs are evolved and validated against a *chunk* spec, exactly as + # the codec pipeline does at run time; see evolve_and_validate_codecs. + chunk_spec = ArraySpec( + shape=representative_chunk_shape(chunk_grid_parsed), dtype=data_type, fill_value=fill_value_parsed, config=ArrayConfig.from_dict({}), # TODO: config is not needed here. prototype=default_buffer_prototype(), # TODO: prototype is not needed here. ) - # Thread the spec through evolution: each codec must be evolved against - # the spec it will actually see at run-time, not the original array spec. - # Earlier array->array codecs may transform the dtype (e.g. cast_value), - # so the spec passed to later codecs must reflect those transformations. - # Per-codec validate() must run before resolve_metadata(), since the - # latter may rely on invariants the former checks (e.g. cast_value - # rejects complex source dtypes that would otherwise crash _do_cast). - evolved: list[Codec] = [] - spec = array_spec - for c in codecs_parsed_partial: - evolved_codec = c.evolve_from_array_spec(spec) - evolved_codec.validate(shape=spec.shape, dtype=spec.dtype, chunk_grid=chunk_grid_parsed) - evolved.append(evolved_codec) - spec = evolved_codec.resolve_metadata(spec) - codecs_parsed = tuple(evolved) + codecs_parsed = evolve_and_validate_codecs( + codecs_parsed_partial, + shape=shape_parsed, + chunk_grid=chunk_grid_parsed, + chunk_spec=chunk_spec, + ) validate_codecs(codecs_parsed_partial, data_type) object.__setattr__(self, "shape", shape_parsed) @@ -557,8 +610,8 @@ def __init__( self._validate_metadata() def _validate_metadata(self) -> None: - if len(self.shape) != self.chunk_grid.ndim: - raise ValueError("`chunk_grid` and `shape` need to have the same number of dimensions.") + # shape/chunk_grid rank agreement is checked in __init__ before the + # codecs are validated, so that a chunk spec of the right rank exists. if isinstance(self.chunk_grid, RectilinearChunkGridMetadata): validate_rectilinear_edges(self.chunk_grid.chunk_shapes, self.shape) if self.dimension_names is not None and len(self.shape) != len(self.dimension_names): @@ -567,8 +620,10 @@ def _validate_metadata(self) -> None: ) if self.fill_value is None: raise ValueError("`fill_value` is required.") - for codec in self.codecs: - codec.validate(shape=self.shape, dtype=self.data_type, chunk_grid=self.chunk_grid) + # Codec validation happens in __init__ (evolve_and_validate_codecs), + # threaded through the chunk spec; re-validating every codec against + # the array-level shape here would wrongly reject chains in which an + # earlier codec changes the chunk's shape or rank. @property def ndim(self) -> int: diff --git a/tests/test_codecs/test_codec_chain_validation.py b/tests/test_codecs/test_codec_chain_validation.py new file mode 100644 index 0000000000..2fee4edbaa --- /dev/null +++ b/tests/test_codecs/test_codec_chain_validation.py @@ -0,0 +1,150 @@ +"""Validation of codec chains in which an earlier array->array codec changes the +shape or rank of a chunk. + +The ``reshape`` extension codec (zarr-extensions) is not implemented in +zarr-python, so a minimal test double is used. Its README explicitly allows +combining ``reshape`` with ``transpose`` to both reorder and reshape; the +``transpose`` order then refers to the *reshaped* rank, so validating it against +the array-level shape must not reject the chain. +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING, Any, Self, cast + +import numpy as np +import pytest + +import zarr +from zarr.abc.codec import ArrayArrayCodec +from zarr.codecs import BytesCodec, ShardingCodec, TransposeCodec +from zarr.core.dtype import Int32 +from zarr.core.metadata.v3 import ArrayV3Metadata, RegularChunkGridMetadata +from zarr.registry import _codec_registries, register_codec + +if TYPE_CHECKING: + from collections.abc import Iterator + + from zarr.core.array_spec import ArraySpec + from zarr.core.buffer import NDBuffer + from zarr.core.common import JSON + + +@dataclass(frozen=True) +class ReshapeCodec(ArrayArrayCodec): + """Minimal stand-in for the zarr-extensions ``reshape`` codec. + + Reshapes every chunk to the explicit ``shape`` (which therefore only makes + sense for a regular chunk grid whose chunks all have the same size). + """ + + shape: tuple[int, ...] + is_fixed_size = True + + @classmethod + def from_dict(cls, data: dict[str, JSON]) -> Self: + config = cast("dict[str, Any]", data["configuration"]) + return cls(shape=tuple(config["shape"])) + + def to_dict(self) -> dict[str, JSON]: + return {"name": "reshape", "configuration": {"shape": list(self.shape)}} + + def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec: + if np.prod(chunk_spec.shape) != np.prod(self.shape): + raise ValueError(f"cannot reshape a chunk of shape {chunk_spec.shape} to {self.shape}") + return replace(chunk_spec, shape=self.shape) + + async def _decode_single(self, chunk_array: NDBuffer, chunk_spec: ArraySpec) -> NDBuffer: + return chunk_array.reshape(chunk_spec.shape) + + async def _encode_single(self, chunk_array: NDBuffer, chunk_spec: ArraySpec) -> NDBuffer: + return chunk_array.reshape(self.shape) + + def compute_encoded_size(self, input_byte_length: int, _chunk_spec: ArraySpec) -> int: + return input_byte_length + + +@pytest.fixture(autouse=True) +def _register_reshape() -> Iterator[None]: + previous = _codec_registries.get("reshape") + register_codec("reshape", ReshapeCodec) + try: + yield + finally: + _codec_registries.pop("reshape", None) + if previous is not None: + _codec_registries["reshape"] = previous + + +SHAPE = (4, 6, 8) +CHUNKS = (2, 3, 4) +# chunk (2, 3, 4) -> (2, 3, 2, 2), then transpose with a rank-4 order +RESHAPE_THEN_TRANSPOSE = (ReshapeCodec(shape=(2, 3, 2, 2)), TransposeCodec(order=(0, 2, 1, 3))) + + +@pytest.mark.parametrize("shards", [None, SHAPE, (2, 6, 8)]) +def test_rank_changing_chain_roundtrip(shards: tuple[int, ...] | None) -> None: + """A reshape+transpose chain is accepted, both standalone and as the inner + codecs of a sharding codec, and round-trips data byte-for-byte.""" + data = np.arange(np.prod(SHAPE), dtype="i4").reshape(SHAPE) + a = zarr.create_array( + {}, + shape=SHAPE, + chunks=CHUNKS, + shards=shards, + dtype="i4", + filters=RESHAPE_THEN_TRANSPOSE, + ) + a[:] = data + assert np.array_equal(a[:], data) + + # The persisted metadata must be re-loadable, i.e. the same validation + # must pass when the codecs come from JSON rather than from instances. + reloaded = zarr.open_array(a.store, mode="r") + assert reloaded.metadata == a.metadata + assert np.array_equal(reloaded[:], data) + + +def _metadata(codecs: tuple[Any, ...], chunk_shape: tuple[int, ...] = CHUNKS) -> ArrayV3Metadata: + return ArrayV3Metadata( + shape=SHAPE, + data_type=Int32(), + chunk_grid=RegularChunkGridMetadata(chunk_shape=chunk_shape), + chunk_key_encoding={"name": "default"}, + fill_value=0, + codecs=codecs, + attributes=None, + dimension_names=None, + ) + + +def test_transpose_validated_against_reshaped_rank() -> None: + """After a rank-changing codec, transpose is validated against the new rank: + an order of the *original* rank is now the invalid one.""" + with pytest.raises(ValueError, match="`order` tuple must have as many entries"): + _metadata((ReshapeCodec(shape=(2, 3, 2, 2)), TransposeCodec(order=(2, 1, 0)), BytesCodec())) + + +def test_reshape_validated_against_chunk_shape() -> None: + """The chunk spec, not the array shape, is threaded through resolve_metadata: + a reshape whose size matches the array but not the chunk is rejected.""" + with pytest.raises(ValueError, match="cannot reshape a chunk of shape"): + _metadata((ReshapeCodec(shape=(4, 6, 8)), BytesCodec())) + + +def test_sharding_inner_chain_is_validated() -> None: + """``ShardingCodec.validate`` validates its inner chain against the inner + chunk shape, threading the spec through rank-changing codecs.""" + grid = RegularChunkGridMetadata(chunk_shape=SHAPE) + ok = ShardingCodec(chunk_shape=CHUNKS, codecs=RESHAPE_THEN_TRANSPOSE) + ok.validate(shape=SHAPE, dtype=Int32(), chunk_grid=grid) + + bad = ShardingCodec( + chunk_shape=CHUNKS, + codecs=(ReshapeCodec(shape=(2, 3, 2, 2)), TransposeCodec(order=(2, 1, 0))), + ) + with pytest.raises(ValueError, match="`order` tuple must have as many entries"): + bad.validate(shape=SHAPE, dtype=Int32(), chunk_grid=grid) + with pytest.raises(ValueError, match="`order` tuple must have as many entries"): + _metadata((bad,), chunk_shape=SHAPE) From 8e324b31a3f740adf4010f1ed179060c8ce1b4a0 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 30 Aug 2026 16:20:49 +0200 Subject: [PATCH 02/11] fix(metadata): validate size-sensitive codecs against every distinct rectilinear chunk shape Review feedback on the threaded-chunk-spec validation: after a codec changes the chunk shape, validating the rest of the chain against a single representative (max-edge) chunk shape is unsound for rectilinear grids -- an inner shard size that divides the largest chunk need not divide the others. Concretely, transpose over a rectilinear grid followed by sharding falsely accepted an inner chunk shape that only divided the largest transposed chunk. The representative was also used to *detect* shape changes, which could miss changes affecting only non-representative chunks. `evolve_and_validate_codecs` now threads every distinct chunk shape of the grid (the cross product of per-dimension distinct edges, capped at 4096 with a ZarrUserWarning on truncation) through `resolve_metadata`, and validates each one individually once any codec has changed a chunk shape. The representative spec remains the single spec used for codec evolution and dtype tracking. Assisted-by: ClaudeCode:claude-fable-5 --- src/zarr/core/metadata/v3.py | 89 +++++++++++++++---- .../test_codec_chain_validation.py | 32 ++++++- 2 files changed, 103 insertions(+), 18 deletions(-) diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index 2bb0f15486..375bb9d262 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -1,6 +1,8 @@ from __future__ import annotations +import itertools import json +import warnings from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass, field, replace from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, TypeGuard, cast @@ -36,7 +38,7 @@ from zarr.core.dtype.common import check_dtype_spec_v3 from zarr.core.json_parse import parse_field from zarr.core.metadata.common import parse_attributes -from zarr.errors import MetadataValidationError, NodeTypeValidationError +from zarr.errors import MetadataValidationError, NodeTypeValidationError, ZarrUserWarning from zarr.registry import get_codec_class if TYPE_CHECKING: @@ -129,15 +131,43 @@ def representative_chunk_shape(chunk_grid: ChunkGridMetadata) -> tuple[int, ...] """A single chunk shape standing in for every chunk of ``chunk_grid``. Regular grids have exactly one chunk shape. Rectilinear grids have many; - the largest edge along each dimension is used, which is enough for the - metadata-time uses of this value (rank checks and threading a chunk spec - through ``Codec.resolve_metadata``). + the largest edge along each dimension is used. This is only suitable where + a single shape is structurally required (rank checks, codec evolution) — + size-sensitive validation must consider every distinct chunk shape, see + ``_distinct_chunk_shapes``. """ if isinstance(chunk_grid, RegularChunkGridMetadata): return chunk_grid.chunk_shape return tuple(s if isinstance(s, int) else max(s) for s in chunk_grid.chunk_shapes) +# Bound on the number of distinct chunk shapes threaded through codec-chain +# validation. A rectilinear grid has prod(distinct edges per dimension) +# distinct chunk shapes, which is unbounded in pathological grids. +_MAX_VALIDATED_CHUNK_SHAPES = 4096 + + +def _distinct_chunk_shapes( + chunk_grid: ChunkGridMetadata, limit: int +) -> tuple[list[tuple[int, ...]], bool]: + """Every distinct chunk shape occurring in ``chunk_grid``, up to ``limit``. + + Returns the shapes and whether the enumeration was truncated at ``limit``. + For a rectilinear grid every combination of per-dimension distinct edges + occurs as an actual chunk shape (each edge along one dimension meets each + edge along every other), so this is the full cross product. + """ + if isinstance(chunk_grid, RegularChunkGridMetadata): + return [chunk_grid.chunk_shape], False + per_dim = ( + (s,) if isinstance(s, int) else tuple(dict.fromkeys(s)) for s in chunk_grid.chunk_shapes + ) + shapes = list(itertools.islice(itertools.product(*per_dim), limit + 1)) + if len(shapes) > limit: + return shapes[:limit], True + return shapes, False + + def evolve_and_validate_codecs( codecs: Iterable[Codec], *, @@ -157,11 +187,19 @@ def evolve_and_validate_codecs( ``shape`` and ``chunk_grid`` are the array-level values passed to ``Codec.validate``. They are handed unchanged to every codec until one - changes the chunk shape; from then on the array-level values are no longer - meaningful for the remaining codecs, so they are replaced by the resolved - chunk shape and a regular grid of that shape (the only shape-related facts - that survive a per-chunk reshape). ``Codec.validate`` implementations only - inspect these for rank and divisibility, so this keeps the checks sound. + changes the shape of any chunk; from then on the array-level values are no + longer meaningful for the remaining codecs. Because ``validate`` checks may + be size-sensitive (sharding divisibility), every *distinct* chunk shape of + the grid is threaded through ``resolve_metadata`` and validated + individually — for a rectilinear grid, a single representative shape would + not be sound: an inner chunk size that divides the largest chunk need not + divide the others. Each threaded shape is presented to ``validate`` as a + regular grid of that shape, the only shape-related facts that survive a + per-chunk transformation. + + ``chunk_spec`` (built from the representative chunk shape) is threaded + separately as the single spec used for codec evolution and dtype tracking, + since evolution must produce one codec chain. Per-codec ``validate`` runs before ``resolve_metadata``, since the latter may rely on invariants the former checks (e.g. ``cast_value`` rejects @@ -169,17 +207,34 @@ def evolve_and_validate_codecs( """ out: list[Codec] = [] spec = chunk_spec - stage_shape = shape - stage_grid = chunk_grid + threaded, truncated = _distinct_chunk_shapes(chunk_grid, _MAX_VALIDATED_CHUNK_SHAPES) + shapes_changed = False for codec in codecs: evolved = codec.evolve_from_array_spec(spec) if evolve else codec - evolved.validate(shape=stage_shape, dtype=spec.dtype, chunk_grid=stage_grid) + if not shapes_changed: + evolved.validate(shape=shape, dtype=spec.dtype, chunk_grid=chunk_grid) + else: + for s in threaded: + evolved.validate( + shape=s, dtype=spec.dtype, chunk_grid=RegularChunkGridMetadata(chunk_shape=s) + ) out.append(evolved) - next_spec = evolved.resolve_metadata(spec) - if next_spec.shape != spec.shape: - stage_shape = next_spec.shape - stage_grid = RegularChunkGridMetadata(chunk_shape=next_spec.shape) - spec = next_spec + resolved = list( + dict.fromkeys(evolved.resolve_metadata(replace(spec, shape=s)).shape for s in threaded) + ) + if resolved != threaded: + shapes_changed = True + if truncated: + warnings.warn( + f"A codec changed the chunk shape of a rectilinear grid with more than " + f"{_MAX_VALIDATED_CHUNK_SHAPES} distinct chunk shapes; codec validation " + "only covered a subset of the chunk shapes.", + category=ZarrUserWarning, + stacklevel=2, + ) + truncated = False + threaded = resolved + spec = evolved.resolve_metadata(spec) return tuple(out) diff --git a/tests/test_codecs/test_codec_chain_validation.py b/tests/test_codecs/test_codec_chain_validation.py index 2fee4edbaa..f7914904a4 100644 --- a/tests/test_codecs/test_codec_chain_validation.py +++ b/tests/test_codecs/test_codec_chain_validation.py @@ -20,7 +20,11 @@ from zarr.abc.codec import ArrayArrayCodec from zarr.codecs import BytesCodec, ShardingCodec, TransposeCodec from zarr.core.dtype import Int32 -from zarr.core.metadata.v3 import ArrayV3Metadata, RegularChunkGridMetadata +from zarr.core.metadata.v3 import ( + ArrayV3Metadata, + RectilinearChunkGridMetadata, + RegularChunkGridMetadata, +) from zarr.registry import _codec_registries, register_codec if TYPE_CHECKING: @@ -148,3 +152,29 @@ def test_sharding_inner_chain_is_validated() -> None: bad.validate(shape=SHAPE, dtype=Int32(), chunk_grid=grid) with pytest.raises(ValueError, match="`order` tuple must have as many entries"): _metadata((bad,), chunk_shape=SHAPE) + + +def _rectilinear_transpose_sharding_metadata(inner: tuple[int, int]) -> ArrayV3Metadata: + """Rectilinear grid (chunks (4,5) and (6,5)), transposed, then sharded.""" + return ArrayV3Metadata( + shape=(10, 5), + data_type=Int32(), + chunk_grid=RectilinearChunkGridMetadata(chunk_shapes=((4, 6), 5)), + chunk_key_encoding={"name": "default"}, + fill_value=0, + codecs=(TransposeCodec(order=(1, 0)), ShardingCodec(chunk_shape=inner)), + attributes=None, + dimension_names=None, + ) + + +def test_rectilinear_every_chunk_shape_validated() -> None: + """Under a rectilinear grid, size-sensitive validation after a + shape-changing codec must consider every distinct chunk shape, not a single + representative: an inner shard size dividing the largest transposed chunk + (5,6) but not the smaller (5,4) is rejected.""" + with zarr.config.set({"array.rectilinear_chunks": True}): + with pytest.raises(ValueError, match="not\\s+divisible"): + _rectilinear_transpose_sharding_metadata((5, 3)) + # an inner shape dividing both transposed chunk shapes is accepted + _rectilinear_transpose_sharding_metadata((5, 2)) From 7c1d01bf118cad24ac97cde208d781ad7f41e958 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 30 Aug 2026 19:15:21 +0200 Subject: [PATCH 03/11] test(codecs): property-based tests for shape-changing codec chain validation Two hypothesis oracles over the threaded-chunk-spec validation: - acceptance implies round-trip: any reshape of a chunk into a valid factorization followed by a transpose of the reshaped rank (with and without sharding) is accepted, encodes/decodes losslessly, and its metadata survives JSON serialization; a transpose order of any other rank is rejected. - transpose-then-shard over a rectilinear grid is accepted exactly when every chunk shape in the grid, transposed, is divisible by the inner shard shape (verified against a brute-force cross-product oracle; this test fails on the max-edge-representative implementation). Assisted-by: ClaudeCode:claude-fable-5 --- .../test_codec_chain_validation_properties.py | 186 ++++++++++++++++++ 1 file changed, 186 insertions(+) create mode 100644 tests/test_codecs/test_codec_chain_validation_properties.py diff --git a/tests/test_codecs/test_codec_chain_validation_properties.py b/tests/test_codecs/test_codec_chain_validation_properties.py new file mode 100644 index 0000000000..d55515feb9 --- /dev/null +++ b/tests/test_codecs/test_codec_chain_validation_properties.py @@ -0,0 +1,186 @@ +"""Property-based tests for codec-chain validation with shape-changing codecs. + +Two invariants are tested against explicit oracles: + +1. Acceptance implies round-trip: any reshape+transpose chain that metadata + validation accepts must encode and decode data losslessly (and its metadata + must survive JSON serialization), while a transpose order of the wrong rank + must be rejected. + +2. For a rectilinear grid followed by a shape-changing codec and a + size-sensitive codec (sharding), acceptance must exactly equal the oracle + "every chunk shape in the grid, transformed by the chain, satisfies the + size constraint" — not just the largest chunk (see + ``evolve_and_validate_codecs``). +""" + +from __future__ import annotations + +import itertools +import math +from typing import TYPE_CHECKING + +import numpy as np +import pytest + +if TYPE_CHECKING: + from collections.abc import Iterator + +import zarr +from zarr.codecs import ShardingCodec, TransposeCodec +from zarr.core.dtype import Int32 +from zarr.core.metadata.v3 import ArrayV3Metadata, RectilinearChunkGridMetadata +from zarr.registry import _codec_registries, register_codec + +from .test_codec_chain_validation import ReshapeCodec + +pytest.importorskip("hypothesis") + +import hypothesis.strategies as st +from hypothesis import given, settings + + +@pytest.fixture(scope="module", autouse=True) +def _register_reshape() -> Iterator[None]: + previous = _codec_registries.get("reshape") + register_codec("reshape", ReshapeCodec) + try: + yield + finally: + _codec_registries.pop("reshape", None) + if previous is not None: + _codec_registries["reshape"] = previous + + +@st.composite +def reshape_transpose_cases( + draw: st.DrawFn, +) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...] | None, tuple[int, ...]]: + """(array shape, chunk shape, shard shape or None, reshape target). + + The reshape target is a valid per-chunk factorization: each chunk dimension + is either kept or split into two factors, so the target always has the same + total size as the chunk but generally a different rank. + """ + ndim = draw(st.integers(min_value=1, max_value=3)) + chunks = tuple(draw(st.integers(min_value=1, max_value=4)) for _ in range(ndim)) + if draw(st.booleans()): + shards = tuple(c * draw(st.integers(min_value=1, max_value=2)) for c in chunks) + else: + shards = None + outer = shards if shards is not None else chunks + shape = tuple(o * draw(st.integers(min_value=1, max_value=2)) for o in outer) + target: list[int] = [] + for c in chunks: + if draw(st.booleans()): + divisor = draw(st.sampled_from([d for d in range(1, c + 1) if c % d == 0])) + target.extend([divisor, c // divisor]) + else: + target.append(c) + return shape, chunks, shards, tuple(target) + + +@settings(deadline=None) +@given(case=reshape_transpose_cases(), data=st.data()) +def test_accepted_reshape_transpose_chain_roundtrips( + case: tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...] | None, tuple[int, ...]], + data: st.DataObject, +) -> None: + """A reshape to any valid chunk factorization, followed by a transpose with + any permutation of the reshaped rank, is accepted and round-trips.""" + shape, chunks, shards, target = case + order = tuple(data.draw(st.permutations(range(len(target))), label="order")) + arr = zarr.create_array( + {}, + shape=shape, + chunks=chunks, + shards=shards, + dtype="i4", + filters=[ReshapeCodec(shape=target), TransposeCodec(order=order)], + ) + expected = np.arange(math.prod(shape), dtype="i4").reshape(shape) + arr[:] = expected + assert np.array_equal(arr[:], expected) + # validation must be stable across JSON serialization + assert ArrayV3Metadata.from_dict(arr.metadata.to_dict()) == arr.metadata + + +@settings(deadline=None) +@given(case=reshape_transpose_cases(), data=st.data()) +def test_wrong_rank_transpose_after_reshape_rejected( + case: tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...] | None, tuple[int, ...]], + data: st.DataObject, +) -> None: + """A transpose order whose rank differs from the reshaped rank is rejected.""" + shape, chunks, shards, target = case + wrong_rank = data.draw( + st.integers(min_value=1, max_value=len(target) + 2).filter(lambda n: n != len(target)), + label="wrong_rank", + ) + order = tuple(data.draw(st.permutations(range(wrong_rank)), label="order")) + with pytest.raises(ValueError, match="order"): + zarr.create_array( + {}, + shape=shape, + chunks=chunks, + shards=shards, + dtype="i4", + filters=[ReshapeCodec(shape=target), TransposeCodec(order=order)], + ) + + +@st.composite +def rectilinear_transpose_sharding_cases( + draw: st.DrawFn, +) -> tuple[tuple[int | tuple[int, ...], ...], tuple[int, ...], tuple[int, ...]]: + """(rectilinear chunk_shapes, transpose order, inner shard shape).""" + ndim = draw(st.integers(min_value=2, max_value=3)) + chunk_shapes: list[int | tuple[int, ...]] = [] + for _ in range(ndim): + edges = draw(st.lists(st.integers(min_value=1, max_value=6), min_size=1, max_size=3)) + # exercise the bare-int (uniform edge) spelling as well + if len(edges) == 1 and draw(st.booleans()): + chunk_shapes.append(edges[0]) + else: + chunk_shapes.append(tuple(edges)) + order = tuple(draw(st.permutations(range(ndim)))) + inner = tuple(draw(st.integers(min_value=1, max_value=6)) for _ in range(ndim)) + return tuple(chunk_shapes), order, inner + + +@settings(deadline=None) +@given(case=rectilinear_transpose_sharding_cases()) +def test_rectilinear_transpose_sharding_matches_oracle( + case: tuple[tuple[int | tuple[int, ...], ...], tuple[int, ...], tuple[int, ...]], +) -> None: + """transpose-then-shard over a rectilinear grid is accepted exactly when + every transposed chunk shape is divisible by the inner shard shape.""" + chunk_shapes, order, inner = case + per_dim = tuple((e,) if isinstance(e, int) else e for e in chunk_shapes) + oracle_ok = all( + all(chunk[order[i]] % inner[i] == 0 for i in range(len(inner))) + for chunk in itertools.product(*per_dim) + ) + # array shape: bare-int (uniform) edges cover any extent; explicit edge + # lists must sum to at least the extent. + shape = tuple(e if isinstance(e, int) else sum(e) for e in chunk_shapes) + + def build() -> ArrayV3Metadata: + return ArrayV3Metadata( + shape=shape, + data_type=Int32(), + chunk_grid=RectilinearChunkGridMetadata(chunk_shapes=chunk_shapes), + chunk_key_encoding={"name": "default"}, + fill_value=0, + codecs=(TransposeCodec(order=order), ShardingCodec(chunk_shape=inner)), + attributes=None, + dimension_names=None, + ) + + with zarr.config.set({"array.rectilinear_chunks": True}): + if oracle_ok: + meta = build() + assert ArrayV3Metadata.from_dict(meta.to_dict()) == meta + else: + with pytest.raises(ValueError, match="not\\s+divisible"): + build() From db9bd2a90c2482b4ee7d6143b059e59694094cfd Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 12 Sep 2026 21:13:07 +0200 Subject: [PATCH 04/11] fix(codecs): validate sharded chains with the real chunk spec Validate inner codecs during evolution, where the actual fill value is available, and cover fill-changing chains with a round-trip property. Assisted-by: Codex:GPT-6 --- changes/+codec-chain-validation.bugfix.md | 2 +- src/zarr/codecs/sharding.py | 28 +++++-------------- .../test_codec_chain_validation.py | 9 +++--- .../test_codec_chain_validation_properties.py | 20 +++++++++++++ 4 files changed, 33 insertions(+), 26 deletions(-) diff --git a/changes/+codec-chain-validation.bugfix.md b/changes/+codec-chain-validation.bugfix.md index 5415125cba..117ed66e6a 100644 --- a/changes/+codec-chain-validation.bugfix.md +++ b/changes/+codec-chain-validation.bugfix.md @@ -1 +1 @@ -Codec chains are now validated against the chunk spec threaded through each codec's `resolve_metadata`, the same way the codec pipeline resolves it at encode time, instead of against the array-level shape. Previously, an `array -> array` codec that changes a chunk's shape or rank (such as the zarr-extensions `reshape` codec) followed by a codec whose configuration refers to the transformed chunk (e.g. `transpose`) was wrongly rejected. The sharding codec now validates its inner codec chain in the same way. +Codec chains are now validated against the chunk spec threaded through each codec's `resolve_metadata`, the same way the codec pipeline resolves it at encode time, instead of against the array-level shape. Previously, an `array -> array` codec that changes a chunk's shape or rank (such as the zarr-extensions `reshape` codec) followed by a codec whose configuration refers to the transformed chunk (e.g. `transpose`) was wrongly rejected. The sharding codec now validates its inner codec chain in the same way. Inner sharding chains are validated during evolution using the actual fill value, so valid fill-changing codecs are not rejected against a fabricated default. diff --git a/src/zarr/codecs/sharding.py b/src/zarr/codecs/sharding.py index 058cf8dc3f..7ede2f6a29 100644 --- a/src/zarr/codecs/sharding.py +++ b/src/zarr/codecs/sharding.py @@ -577,10 +577,13 @@ def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self: ------- This codec with the evolved code chain. """ - from zarr.core.chunk_utils import evolve_codecs - - shard_spec = self._get_chunk_spec(array_spec) - evolved_codecs = evolve_codecs(self.codecs, shard_spec) + chunk_spec = self._get_chunk_spec(array_spec) + evolved_codecs = evolve_and_validate_codecs( + self.codecs, + shape=self.chunk_shape, + chunk_grid=RegularChunkGridMetadata(chunk_shape=self.chunk_shape), + chunk_spec=chunk_spec, + ) if evolved_codecs != self.codecs: return replace(self, codecs=evolved_codecs) return self @@ -614,23 +617,6 @@ def validate( f"Chunk edge length {edge} in dimension {i} is not " f"divisible by the shard's inner chunk size {inner}." ) - # The inner codecs see chunks of `self.chunk_shape`; validate them - # against that, threading the chunk spec through the chain exactly as - # the top-level metadata does (an inner reshape may change the rank - # seen by a following transpose). - evolve_and_validate_codecs( - self.codecs, - shape=self.chunk_shape, - chunk_grid=RegularChunkGridMetadata(chunk_shape=self.chunk_shape), - chunk_spec=ArraySpec( - shape=self.chunk_shape, - dtype=dtype, - fill_value=dtype.default_scalar(), - config=ArrayConfig.from_dict({}), - prototype=default_buffer_prototype(), - ), - evolve=False, - ) def _get_inner_chunk_transform(self, shard_spec: ArraySpec) -> Any: """The synchronous transform for the inner codec chain. diff --git a/tests/test_codecs/test_codec_chain_validation.py b/tests/test_codecs/test_codec_chain_validation.py index f7914904a4..fcbdcb6c87 100644 --- a/tests/test_codecs/test_codec_chain_validation.py +++ b/tests/test_codecs/test_codec_chain_validation.py @@ -138,8 +138,11 @@ def test_reshape_validated_against_chunk_shape() -> None: def test_sharding_inner_chain_is_validated() -> None: - """``ShardingCodec.validate`` validates its inner chain against the inner - chunk shape, threading the spec through rank-changing codecs.""" + """Metadata construction validates inner codecs with the real chunk spec. + + Direct ``validate`` only has geometry and dtype, not the fill value needed + to resolve arbitrary inner codecs; evolution supplies that full context. + """ grid = RegularChunkGridMetadata(chunk_shape=SHAPE) ok = ShardingCodec(chunk_shape=CHUNKS, codecs=RESHAPE_THEN_TRANSPOSE) ok.validate(shape=SHAPE, dtype=Int32(), chunk_grid=grid) @@ -148,8 +151,6 @@ def test_sharding_inner_chain_is_validated() -> None: chunk_shape=CHUNKS, codecs=(ReshapeCodec(shape=(2, 3, 2, 2)), TransposeCodec(order=(2, 1, 0))), ) - with pytest.raises(ValueError, match="`order` tuple must have as many entries"): - bad.validate(shape=SHAPE, dtype=Int32(), chunk_grid=grid) with pytest.raises(ValueError, match="`order` tuple must have as many entries"): _metadata((bad,), chunk_shape=SHAPE) diff --git a/tests/test_codecs/test_codec_chain_validation_properties.py b/tests/test_codecs/test_codec_chain_validation_properties.py index d55515feb9..4f839ed286 100644 --- a/tests/test_codecs/test_codec_chain_validation_properties.py +++ b/tests/test_codecs/test_codec_chain_validation_properties.py @@ -184,3 +184,23 @@ def build() -> ArrayV3Metadata: else: with pytest.raises(ValueError, match="not\\s+divisible"): build() + + +@given(offset=st.integers(min_value=1, max_value=254), sharded=st.booleans()) +def test_inner_validation_uses_the_actual_fill_value(offset: int, sharded: bool) -> None: + """Shard validation must not resolve a fill-changing codec against a made-up zero.""" + from zarr.codecs.scale_offset import ScaleOffset + + array = zarr.create_array( + {}, + shape=(8,), + chunks=(2,), + shards=(4,) if sharded else None, + dtype="u1", + fill_value=offset, + filters=[ScaleOffset(offset=offset)], + ) + array[:] = offset + 1 + assert np.array_equal(array[:], np.full((8,), offset + 1, dtype="u1")) + reloaded = zarr.open_array(array.store, mode="r") + assert np.array_equal(reloaded[:], array[:]) From aaab1937f5ddcde3f682e2b7b2c1086bae9b6a12 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 12 Sep 2026 21:25:38 +0200 Subject: [PATCH 05/11] fix(metadata): name the failing codec in chain validation errors Review fixes for the threaded codec-chain validation: - Attach a note to any exception raised by a codec's evolve_from_array_spec, validate or resolve_metadata naming the codec's position in the chain, its class and the shape it was checked against. Codec messages talk about "the array", which is misleading once an earlier codec has changed the chunk shape; the exception type and message are unchanged. - Drop the unused `evolve` parameter of evolve_and_validate_codecs left over from the removed ShardingCodec.validate inner-chain check. - Document in ShardingCodec.evolve_from_array_spec why the inner chain is validated there (validate has no fill value) and against which grid. - Docstrings use single backticks; the towncrier fragment is 303.bugfix.md. - Tests: one parametrized happy-path test for rectilinear grids plus one test per error case; the sharding inner-chain test asserts the note. Assisted-by: ClaudeCode:claude-fable-5-1 Co-Authored-By: Claude Fable 5.1 --- ...ain-validation.bugfix.md => 303.bugfix.md} | 2 +- src/zarr/codecs/sharding.py | 9 +- src/zarr/core/metadata/v3.py | 106 ++++++++++++------ .../test_codec_chain_validation.py | 39 ++++--- 4 files changed, 103 insertions(+), 53 deletions(-) rename changes/{+codec-chain-validation.bugfix.md => 303.bugfix.md} (77%) diff --git a/changes/+codec-chain-validation.bugfix.md b/changes/303.bugfix.md similarity index 77% rename from changes/+codec-chain-validation.bugfix.md rename to changes/303.bugfix.md index 117ed66e6a..8dab311fba 100644 --- a/changes/+codec-chain-validation.bugfix.md +++ b/changes/303.bugfix.md @@ -1 +1 @@ -Codec chains are now validated against the chunk spec threaded through each codec's `resolve_metadata`, the same way the codec pipeline resolves it at encode time, instead of against the array-level shape. Previously, an `array -> array` codec that changes a chunk's shape or rank (such as the zarr-extensions `reshape` codec) followed by a codec whose configuration refers to the transformed chunk (e.g. `transpose`) was wrongly rejected. The sharding codec now validates its inner codec chain in the same way. Inner sharding chains are validated during evolution using the actual fill value, so valid fill-changing codecs are not rejected against a fabricated default. +Codec chains are now validated against the chunk spec threaded through each codec's `resolve_metadata`, the same way the codec pipeline resolves it at encode time, instead of against the array-level shape. Previously, an `array -> array` codec that changes a chunk's shape or rank (such as the zarr-extensions `reshape` codec) followed by a codec whose configuration refers to the transformed chunk (e.g. `transpose`) was wrongly rejected. The sharding codec now validates its inner codec chain in the same way. Inner sharding chains are validated during evolution using the actual fill value, so valid fill-changing codecs are not rejected against a fabricated default. An error raised while evolving or validating a codec now carries a note naming the codec's position in the chain and the shape it was checked against. diff --git a/src/zarr/codecs/sharding.py b/src/zarr/codecs/sharding.py index 7ede2f6a29..46854e63e3 100644 --- a/src/zarr/codecs/sharding.py +++ b/src/zarr/codecs/sharding.py @@ -560,7 +560,7 @@ def to_dict(self) -> dict[str, JSON]: } def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self: - """Thread the spec through the inner chain. + """Thread the spec through the inner chain, evolving and validating it. Each codec is evolved against the spec produced by the previous one. Evolving every codec against the same unthreaded spec is the bug shape that @@ -568,6 +568,13 @@ def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self: method runs on the real array-creation path, baking the damaged chain into the evolved instance before the transform builders ever run. + The inner chain is validated here rather than in `validate`, because + only the array spec carries the fill value some inner codecs need to + resolve their metadata (e.g. `scale_offset`); `validate` only has the + geometry and dtype. Inner codecs see chunks of `chunk_shape`, so that + is the shape and (regular) grid they are validated against, threaded + through any shape-changing inner codec as at the top level. + Parameters ---------- array_spec diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index 375bb9d262..ac1864099a 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -128,13 +128,13 @@ def validate_codecs(codecs: tuple[Codec, ...], dtype: ZDType[TBaseDType, TBaseSc def representative_chunk_shape(chunk_grid: ChunkGridMetadata) -> tuple[int, ...]: - """A single chunk shape standing in for every chunk of ``chunk_grid``. + """A single chunk shape standing in for every chunk of `chunk_grid`. Regular grids have exactly one chunk shape. Rectilinear grids have many; the largest edge along each dimension is used. This is only suitable where - a single shape is structurally required (rank checks, codec evolution) — + a single shape is structurally required (rank checks, codec evolution); size-sensitive validation must consider every distinct chunk shape, see - ``_distinct_chunk_shapes``. + `_distinct_chunk_shapes`. """ if isinstance(chunk_grid, RegularChunkGridMetadata): return chunk_grid.chunk_shape @@ -150,9 +150,9 @@ def representative_chunk_shape(chunk_grid: ChunkGridMetadata) -> tuple[int, ...] def _distinct_chunk_shapes( chunk_grid: ChunkGridMetadata, limit: int ) -> tuple[list[tuple[int, ...]], bool]: - """Every distinct chunk shape occurring in ``chunk_grid``, up to ``limit``. + """Every distinct chunk shape occurring in `chunk_grid`, up to `limit`. - Returns the shapes and whether the enumeration was truncated at ``limit``. + Returns the shapes and whether the enumeration was truncated at `limit`. For a rectilinear grid every combination of per-dimension distinct edges occurs as an actual chunk shape (each edge along one dimension meets each edge along every other), so this is the full cross product. @@ -168,60 +168,94 @@ def _distinct_chunk_shapes( return shapes, False +def _note_codec(exc: BaseException, position: int, codec: Codec, shape: tuple[int, ...]) -> None: + """Attach a note naming the codec and the shape it was checked against. + + Codec error messages talk about "the array", but after a shape-changing + codec they describe the transformed chunk; the note makes that visible in + the traceback without changing the exception's type or message. + """ + exc.add_note( + f"Raised by codec {position} of the chain, {type(codec).__name__}, checked against " + f"shape {shape}." + ) + + def evolve_and_validate_codecs( codecs: Iterable[Codec], *, shape: tuple[int, ...], chunk_grid: ChunkGridMetadata, chunk_spec: ArraySpec, - evolve: bool = True, ) -> tuple[Codec, ...]: - """Evolve (optionally) and validate a codec chain, threading the chunk spec. + """Evolve and validate a codec chain, threading the chunk spec. Each codec is evolved and validated against the chunk spec produced by the - previous codec's ``resolve_metadata`` — the same spec it will see at - encode/decode time — not against the array-level metadata. Earlier - array->array codecs may change the dtype (``cast_value``) or the shape and - even the rank of a chunk (the ``reshape`` extension codec, which the spec - explicitly allows to be followed by ``transpose``). - - ``shape`` and ``chunk_grid`` are the array-level values passed to - ``Codec.validate``. They are handed unchanged to every codec until one + previous codec's `resolve_metadata`, the same spec it will see at + encode/decode time, not against the array-level metadata. Earlier + array->array codecs may change the dtype (`cast_value`) or the shape and + even the rank of a chunk (the `reshape` extension codec, which the spec + explicitly allows to be followed by `transpose`). + + `shape` and `chunk_grid` are the array-level values passed to + `Codec.validate`. They are handed unchanged to every codec until one changes the shape of any chunk; from then on the array-level values are no - longer meaningful for the remaining codecs. Because ``validate`` checks may + longer meaningful for the remaining codecs. Because `validate` checks may be size-sensitive (sharding divisibility), every *distinct* chunk shape of - the grid is threaded through ``resolve_metadata`` and validated - individually — for a rectilinear grid, a single representative shape would + the grid is threaded through `resolve_metadata` and validated + individually; for a rectilinear grid, a single representative shape would not be sound: an inner chunk size that divides the largest chunk need not - divide the others. Each threaded shape is presented to ``validate`` as a + divide the others. Each threaded shape is presented to `validate` as a regular grid of that shape, the only shape-related facts that survive a per-chunk transformation. - ``chunk_spec`` (built from the representative chunk shape) is threaded + `chunk_spec` (built from the representative chunk shape) is threaded separately as the single spec used for codec evolution and dtype tracking, since evolution must produce one codec chain. - Per-codec ``validate`` runs before ``resolve_metadata``, since the latter - may rely on invariants the former checks (e.g. ``cast_value`` rejects - complex source dtypes that would otherwise crash ``_do_cast``). + Per-codec `validate` runs before `resolve_metadata`, since the latter + may rely on invariants the former checks (e.g. `cast_value` rejects + complex source dtypes that would otherwise crash `_do_cast`). + + Any exception raised by a codec's `evolve_from_array_spec`, `validate` or + `resolve_metadata` is re-raised unchanged with a note naming the codec's + position in the chain and the shape it was checked against. """ out: list[Codec] = [] spec = chunk_spec threaded, truncated = _distinct_chunk_shapes(chunk_grid, _MAX_VALIDATED_CHUNK_SHAPES) shapes_changed = False - for codec in codecs: - evolved = codec.evolve_from_array_spec(spec) if evolve else codec - if not shapes_changed: - evolved.validate(shape=shape, dtype=spec.dtype, chunk_grid=chunk_grid) - else: - for s in threaded: - evolved.validate( - shape=s, dtype=spec.dtype, chunk_grid=RegularChunkGridMetadata(chunk_shape=s) - ) - out.append(evolved) - resolved = list( - dict.fromkeys(evolved.resolve_metadata(replace(spec, shape=s)).shape for s in threaded) + for position, codec in enumerate(codecs): + # (shape, chunk_grid) pairs handed to validate: the array-level pair + # until a codec changes chunk shapes, then each distinct chunk shape + # as its own regular grid. + stages = ( + [(s, RegularChunkGridMetadata(chunk_shape=s)) for s in threaded] + if shapes_changed + else [(shape, chunk_grid)] ) + try: + evolved = codec.evolve_from_array_spec(spec) + except Exception as e: + _note_codec(e, position, codec, spec.shape) + raise + for stage_shape, stage_grid in stages: + try: + evolved.validate(shape=stage_shape, dtype=spec.dtype, chunk_grid=stage_grid) + except Exception as e: + _note_codec(e, position, codec, stage_shape) + raise + out.append(evolved) + try: + resolved = list( + dict.fromkeys( + evolved.resolve_metadata(replace(spec, shape=s)).shape for s in threaded + ) + ) + next_spec = evolved.resolve_metadata(spec) + except Exception as e: + _note_codec(e, position, codec, spec.shape) + raise if resolved != threaded: shapes_changed = True if truncated: @@ -234,7 +268,7 @@ def evolve_and_validate_codecs( ) truncated = False threaded = resolved - spec = evolved.resolve_metadata(spec) + spec = next_spec return tuple(out) diff --git a/tests/test_codecs/test_codec_chain_validation.py b/tests/test_codecs/test_codec_chain_validation.py index fcbdcb6c87..2327398c7e 100644 --- a/tests/test_codecs/test_codec_chain_validation.py +++ b/tests/test_codecs/test_codec_chain_validation.py @@ -138,21 +138,19 @@ def test_reshape_validated_against_chunk_shape() -> None: def test_sharding_inner_chain_is_validated() -> None: - """Metadata construction validates inner codecs with the real chunk spec. - - Direct ``validate`` only has geometry and dtype, not the fill value needed - to resolve arbitrary inner codecs; evolution supplies that full context. - """ - grid = RegularChunkGridMetadata(chunk_shape=SHAPE) - ok = ShardingCodec(chunk_shape=CHUNKS, codecs=RESHAPE_THEN_TRANSPOSE) - ok.validate(shape=SHAPE, dtype=Int32(), chunk_grid=grid) - + """Metadata construction validates a sharding codec's inner chain against + the inner chunk spec (the accepted case is `test_rank_changing_chain_roundtrip` + with `shards`). Direct `ShardingCodec.validate` only has geometry and dtype, + not the fill value needed to resolve arbitrary inner codecs, so the inner + chain is validated during evolution, which has the full spec.""" bad = ShardingCodec( chunk_shape=CHUNKS, codecs=(ReshapeCodec(shape=(2, 3, 2, 2)), TransposeCodec(order=(2, 1, 0))), ) - with pytest.raises(ValueError, match="`order` tuple must have as many entries"): + with pytest.raises(ValueError, match="`order` tuple must have as many entries") as exc_info: _metadata((bad,), chunk_shape=SHAPE) + # the error names the codec that raised and the shape it was checked against + assert any("TransposeCodec" in n and "(2, 3, 2, 2)" in n for n in exc_info.value.__notes__) def _rectilinear_transpose_sharding_metadata(inner: tuple[int, int]) -> ArrayV3Metadata: @@ -169,13 +167,24 @@ def _rectilinear_transpose_sharding_metadata(inner: tuple[int, int]) -> ArrayV3M ) +@pytest.mark.parametrize("inner", [(5, 2), (5, 1), (1, 2)]) +def test_rectilinear_shape_change_accepts_inner_dividing_every_chunk( + inner: tuple[int, int], +) -> None: + """After a shape-changing codec on a rectilinear grid, an inner shard shape + dividing every transposed chunk shape ((5,4) and (5,6)) is accepted.""" + with zarr.config.set({"array.rectilinear_chunks": True}): + meta = _rectilinear_transpose_sharding_metadata(inner) + assert ArrayV3Metadata.from_dict(meta.to_dict()) == meta + + def test_rectilinear_every_chunk_shape_validated() -> None: """Under a rectilinear grid, size-sensitive validation after a shape-changing codec must consider every distinct chunk shape, not a single representative: an inner shard size dividing the largest transposed chunk (5,6) but not the smaller (5,4) is rejected.""" - with zarr.config.set({"array.rectilinear_chunks": True}): - with pytest.raises(ValueError, match="not\\s+divisible"): - _rectilinear_transpose_sharding_metadata((5, 3)) - # an inner shape dividing both transposed chunk shapes is accepted - _rectilinear_transpose_sharding_metadata((5, 2)) + with ( + zarr.config.set({"array.rectilinear_chunks": True}), + pytest.raises(ValueError, match="not\\s+divisible"), + ): + _rectilinear_transpose_sharding_metadata((5, 3)) From d849cf39e2ece213e468ff50da80539954868aa8 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 12 Sep 2026 21:26:18 +0200 Subject: [PATCH 06/11] test(properties): transpose ahead of sharding is accepted iff it tiles the transposed chunk Add the `transposed_sharding_chains` strategy: a codec chain with a TransposeCodec (random permutation) ahead of a ShardingCodec in one of three layouts (transpose then shard; nested shard with the transpose between the levels; transpose inside a shard as the always-valid control), together with an oracle for its validity (every edge of the sharding codec's chunk shape divides the transposed edge it applies to). Half the drawn chains are invalid, breaking exactly one axis. The property asserts that create_array accepts a chain exactly when the oracle says it is valid, and that an accepted chain round-trips its data and its persisted metadata. On main before the fix it fails two ways: a nested sharding codec's inner chain is never validated, so an invalid inner chunk shape is accepted and reads back wrong data, and a valid transpose-then-shard chain is rejected because the sharding codec was validated against the untransposed chunk grid. Assisted-by: ClaudeCode:claude-fable-5-1 Co-Authored-By: Claude Fable 5.1 --- src/zarr/testing/strategies.py | 91 ++++++++++++++++++++++++++++++++++ tests/test_properties.py | 49 ++++++++++++++++++ 2 files changed, 140 insertions(+) diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index db01697f1e..620b982eb2 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -23,6 +23,7 @@ from zarr.codecs.bytes import BytesCodec from zarr.codecs.crc32c_ import Crc32cCodec from zarr.codecs.sharding import SUBCHUNK_WRITE_ORDER, ShardingCodec, SubchunkWriteOrder +from zarr.codecs.transpose import TransposeCodec from zarr.codecs.zstd import ZstdCodec from zarr.core.array import Array, CompressorsLike, SerializerLike from zarr.core.chunk_key_encodings import DefaultChunkKeyEncoding @@ -280,6 +281,96 @@ def _sharding_codecs( ) +def _divisors(n: int) -> st.SearchStrategy[int]: + return st.sampled_from([d for d in range(1, n + 1) if n % d == 0]) + + +@st.composite +def transposed_sharding_chains( + draw: st.DrawFn, *, shape: tuple[int, ...] +) -> tuple[tuple[int, ...], list[Codec] | None, ShardingCodec, bool]: + """A codec chain with a `TransposeCodec` ahead of a `ShardingCodec`, and its validity. + + Returns `(chunks, filters, serializer, valid)` for + `zarr.create_array(shape=shape, chunks=chunks, filters=filters, serializer=serializer)`. + One of three layouts is drawn: + + - `transpose then shard`: `filters=[transpose]`, the serializer shards the + transposed chunk. + - `nested shard, transpose between levels`: an outer `ShardingCodec` whose + inner chain is `[transpose, ShardingCodec]`. + - `transpose inside shard`: a `ShardingCodec` whose inner chain is + `[transpose, BytesCodec]`; a control that is always valid. + + A sharding codec placed after the transpose sees transposed chunks, so its + `chunk_shape` must divide the transposed edges, not the array's chunk + grid. That inner chunk shape is drawn without regard to divisibility, so + about half of the drawn chains are invalid; `valid` is the oracle: every + edge of the sharding codec's chunk shape divides the transposed edge it + applies to. Any shard shape drawn for the outer codec always divides the + array's chunk grid, so `valid` is decided by the transpose alone. + """ + ndim = len(shape) + # Edges are drawn uniformly rather than via ``chunk_shapes`` (which favors + # many chunks of edge 1): unequal edges are what make a transpose change + # the chunk shape, which is the case this strategy exists for. + chunks = tuple(draw(st.integers(min_value=1, max_value=s)) for s in shape) + order = tuple(draw(st.permutations(range(ndim)), label="transpose order")) + transpose = TransposeCodec(order=order) + layout = draw( + st.sampled_from( + [ + "transpose then shard", + "nested shard, transpose between levels", + "transpose inside shard", + ] + ), + label="layout", + ) + event("transposed sharding chain layout", layout) + + def transposed(edges: tuple[int, ...]) -> tuple[int, ...]: + return tuple(edges[order[i]] for i in range(ndim)) + + event("transpose changes the chunk shape", "yes" if transposed(chunks) != chunks else "no") + + def draw_inner(seen: tuple[int, ...]) -> tuple[tuple[int, ...], bool]: + # Half the chains are valid (every edge a divisor of the transposed + # edge); the other half break exactly one axis with a non-divisor, of + # which `t + 1` guarantees at least one exists. + valid = draw(st.booleans(), label="valid chain") + broken = None if valid else draw(st.integers(min_value=0, max_value=ndim - 1)) + inner = tuple( + draw( + _divisors(t) + if axis != broken + else st.sampled_from([i for i in range(1, t + 2) if t % i != 0]) + ) + for axis, t in enumerate(seen) + ) + assert valid == all(t % i == 0 for t, i in zip(seen, inner, strict=True)) + return inner, valid + + filters: list[Codec] | None = None + if layout == "transpose then shard": + inner, valid = draw_inner(transposed(chunks)) + filters = [transpose] + serializer = ShardingCodec(chunk_shape=inner, codecs=[BytesCodec()]) + else: + mid = tuple(draw(_divisors(c)) for c in chunks) + if layout == "nested shard, transpose between levels": + inner, valid = draw_inner(transposed(mid)) + serializer = ShardingCodec( + chunk_shape=mid, + codecs=[transpose, ShardingCodec(chunk_shape=inner, codecs=[BytesCodec()])], + ) + else: + valid = True + serializer = ShardingCodec(chunk_shape=mid, codecs=[transpose, BytesCodec()]) + event("transposed sharding chain", "valid" if valid else "invalid") + return chunks, filters, serializer, valid + + @st.composite def np_array_and_chunks( draw: st.DrawFn, diff --git a/tests/test_properties.py b/tests/test_properties.py index 2794ad3cb0..5c3f480f46 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -34,6 +34,7 @@ sharded_arrays, simple_arrays, stores, + transposed_sharding_chains, zarr_formats, ) @@ -496,3 +497,51 @@ def test_chunks_param_from_rectilinear_bare_int_roundtrip() -> None: dtype="uint8", ) assert dst.metadata.chunk_grid == grid # type: ignore[union-attr] + + +@given(data=st.data()) +@pytest.mark.filterwarnings( + "ignore:Combining a `sharding_indexed` codec:zarr.errors.ZarrUserWarning" +) +def test_transposed_sharding_chain_validation(data: st.DataObject) -> None: + """A chain with a transpose ahead of a sharding codec is accepted exactly + when the sharding codec's chunk shape divides the *transposed* chunk, and an + accepted chain round-trips its data and its persisted metadata. + + Every codec must be validated against the chunk spec it sees at encode + time. Validating the sharding codec against the array's untransposed chunk + grid instead accepts chains whose inner chunks do not tile the transposed + chunk; such an array is created without error and then silently reads back + wrong data. + """ + from zarr.storage import MemoryStore + + # At least two axes, so that a permuted order can change the chunk shape. + shape = data.draw( + npst.array_shapes(min_dims=2, max_dims=3, min_side=1, max_side=8), label="shape" + ) + chunks, filters, serializer, valid = data.draw(transposed_sharding_chains(shape=shape)) + dtype = data.draw(st.sampled_from([np.dtype("uint8"), np.dtype("int32"), np.dtype("float64")])) + nparray = data.draw(numpy_arrays(shapes=st.just(shape), dtype=dtype), label="array data") + + def build() -> zarr.Array: + return zarr.create_array( + MemoryStore(), + shape=shape, + chunks=chunks, + dtype=dtype, + filters=filters, + serializer=serializer, + compressors=None, + ) + + if not valid: + with pytest.raises(ValueError, match="not divisible"): + build() + return + zarray = build() + zarray[:] = nparray + assert_array_equal(zarray[:], nparray) + reopened = zarr.open_array(zarray.store, mode="r") + assert reopened.metadata == zarray.metadata + assert_array_equal(reopened[:], nparray) From 04dd144c12c310a1197cbf279b53f1de9128b983 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 12 Sep 2026 22:31:27 +0200 Subject: [PATCH 07/11] refactor(metadata): cut the partial-validation path and error notes from chain validation Apply the second-lens review cuts: - Enumerate every distinct rectilinear chunk shape without a cap. The warn-and-continue path validated only a subset of shapes and then accepted the metadata, which is worse than either failing or checking everything. - Drop the exception notes and the try/except blocks around evolve, validate and resolve_metadata; they were not part of the fix. - Move `transposed_sharding_chains` out of the public `zarr.testing.strategies` module into `tests/test_properties.py`, its only user. - Drop the metadata-only rectilinear oracle property; the create_array plus round-trip property in `test_properties.py` covers the same acceptance oracle for regular and nested sharding. Assisted-by: ClaudeCode:claude-fable-5-1 Co-Authored-By: Claude Fable 5.1 --- changes/303.bugfix.md | 2 +- src/zarr/core/metadata/v3.py | 98 +++++-------------- src/zarr/testing/strategies.py | 91 ----------------- .../test_codec_chain_validation.py | 4 +- .../test_codec_chain_validation_properties.py | 81 ++------------- tests/test_properties.py | 93 +++++++++++++++++- 6 files changed, 124 insertions(+), 245 deletions(-) diff --git a/changes/303.bugfix.md b/changes/303.bugfix.md index 8dab311fba..01304976c0 100644 --- a/changes/303.bugfix.md +++ b/changes/303.bugfix.md @@ -1 +1 @@ -Codec chains are now validated against the chunk spec threaded through each codec's `resolve_metadata`, the same way the codec pipeline resolves it at encode time, instead of against the array-level shape. Previously, an `array -> array` codec that changes a chunk's shape or rank (such as the zarr-extensions `reshape` codec) followed by a codec whose configuration refers to the transformed chunk (e.g. `transpose`) was wrongly rejected. The sharding codec now validates its inner codec chain in the same way. Inner sharding chains are validated during evolution using the actual fill value, so valid fill-changing codecs are not rejected against a fabricated default. An error raised while evolving or validating a codec now carries a note naming the codec's position in the chain and the shape it was checked against. +Codec chains are now validated against the chunk spec threaded through each codec's `resolve_metadata`, the same way the codec pipeline resolves it at encode time, instead of against the array-level shape. Previously, an `array -> array` codec that changes a chunk's shape or rank (such as the zarr-extensions `reshape` codec) followed by a codec whose configuration refers to the transformed chunk (e.g. `transpose`) was wrongly rejected. The sharding codec now validates its inner codec chain in the same way. Inner sharding chains are validated during evolution using the actual fill value, so valid fill-changing codecs are not rejected against a fabricated default. \ No newline at end of file diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index ac1864099a..a4767e309c 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -2,7 +2,6 @@ import itertools import json -import warnings from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass, field, replace from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, TypeGuard, cast @@ -38,7 +37,7 @@ from zarr.core.dtype.common import check_dtype_spec_v3 from zarr.core.json_parse import parse_field from zarr.core.metadata.common import parse_attributes -from zarr.errors import MetadataValidationError, NodeTypeValidationError, ZarrUserWarning +from zarr.errors import MetadataValidationError, NodeTypeValidationError from zarr.registry import get_codec_class if TYPE_CHECKING: @@ -141,44 +140,19 @@ def representative_chunk_shape(chunk_grid: ChunkGridMetadata) -> tuple[int, ...] return tuple(s if isinstance(s, int) else max(s) for s in chunk_grid.chunk_shapes) -# Bound on the number of distinct chunk shapes threaded through codec-chain -# validation. A rectilinear grid has prod(distinct edges per dimension) -# distinct chunk shapes, which is unbounded in pathological grids. -_MAX_VALIDATED_CHUNK_SHAPES = 4096 +def _distinct_chunk_shapes(chunk_grid: ChunkGridMetadata) -> list[tuple[int, ...]]: + """Every distinct chunk shape occurring in `chunk_grid`. - -def _distinct_chunk_shapes( - chunk_grid: ChunkGridMetadata, limit: int -) -> tuple[list[tuple[int, ...]], bool]: - """Every distinct chunk shape occurring in `chunk_grid`, up to `limit`. - - Returns the shapes and whether the enumeration was truncated at `limit`. For a rectilinear grid every combination of per-dimension distinct edges occurs as an actual chunk shape (each edge along one dimension meets each edge along every other), so this is the full cross product. """ if isinstance(chunk_grid, RegularChunkGridMetadata): - return [chunk_grid.chunk_shape], False + return [chunk_grid.chunk_shape] per_dim = ( (s,) if isinstance(s, int) else tuple(dict.fromkeys(s)) for s in chunk_grid.chunk_shapes ) - shapes = list(itertools.islice(itertools.product(*per_dim), limit + 1)) - if len(shapes) > limit: - return shapes[:limit], True - return shapes, False - - -def _note_codec(exc: BaseException, position: int, codec: Codec, shape: tuple[int, ...]) -> None: - """Attach a note naming the codec and the shape it was checked against. - - Codec error messages talk about "the array", but after a shape-changing - codec they describe the transformed chunk; the note makes that visible in - the traceback without changing the exception's type or message. - """ - exc.add_note( - f"Raised by codec {position} of the chain, {type(codec).__name__}, checked against " - f"shape {shape}." - ) + return list(itertools.product(*per_dim)) def evolve_and_validate_codecs( @@ -216,59 +190,31 @@ def evolve_and_validate_codecs( Per-codec `validate` runs before `resolve_metadata`, since the latter may rely on invariants the former checks (e.g. `cast_value` rejects complex source dtypes that would otherwise crash `_do_cast`). - - Any exception raised by a codec's `evolve_from_array_spec`, `validate` or - `resolve_metadata` is re-raised unchanged with a note naming the codec's - position in the chain and the shape it was checked against. """ out: list[Codec] = [] spec = chunk_spec - threaded, truncated = _distinct_chunk_shapes(chunk_grid, _MAX_VALIDATED_CHUNK_SHAPES) + threaded = _distinct_chunk_shapes(chunk_grid) shapes_changed = False - for position, codec in enumerate(codecs): - # (shape, chunk_grid) pairs handed to validate: the array-level pair - # until a codec changes chunk shapes, then each distinct chunk shape - # as its own regular grid. - stages = ( - [(s, RegularChunkGridMetadata(chunk_shape=s)) for s in threaded] - if shapes_changed - else [(shape, chunk_grid)] - ) - try: - evolved = codec.evolve_from_array_spec(spec) - except Exception as e: - _note_codec(e, position, codec, spec.shape) - raise - for stage_shape, stage_grid in stages: - try: - evolved.validate(shape=stage_shape, dtype=spec.dtype, chunk_grid=stage_grid) - except Exception as e: - _note_codec(e, position, codec, stage_shape) - raise - out.append(evolved) - try: - resolved = list( - dict.fromkeys( - evolved.resolve_metadata(replace(spec, shape=s)).shape for s in threaded + for codec in codecs: + evolved = codec.evolve_from_array_spec(spec) + # The array-level shape and grid are handed to validate until a codec + # changes chunk shapes; from then on each distinct chunk shape is + # validated as its own regular grid. + if shapes_changed: + for s in threaded: + evolved.validate( + shape=s, dtype=spec.dtype, chunk_grid=RegularChunkGridMetadata(chunk_shape=s) ) - ) - next_spec = evolved.resolve_metadata(spec) - except Exception as e: - _note_codec(e, position, codec, spec.shape) - raise + else: + evolved.validate(shape=shape, dtype=spec.dtype, chunk_grid=chunk_grid) + out.append(evolved) + resolved = list( + dict.fromkeys(evolved.resolve_metadata(replace(spec, shape=s)).shape for s in threaded) + ) if resolved != threaded: shapes_changed = True - if truncated: - warnings.warn( - f"A codec changed the chunk shape of a rectilinear grid with more than " - f"{_MAX_VALIDATED_CHUNK_SHAPES} distinct chunk shapes; codec validation " - "only covered a subset of the chunk shapes.", - category=ZarrUserWarning, - stacklevel=2, - ) - truncated = False threaded = resolved - spec = next_spec + spec = evolved.resolve_metadata(spec) return tuple(out) diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index 620b982eb2..db01697f1e 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -23,7 +23,6 @@ from zarr.codecs.bytes import BytesCodec from zarr.codecs.crc32c_ import Crc32cCodec from zarr.codecs.sharding import SUBCHUNK_WRITE_ORDER, ShardingCodec, SubchunkWriteOrder -from zarr.codecs.transpose import TransposeCodec from zarr.codecs.zstd import ZstdCodec from zarr.core.array import Array, CompressorsLike, SerializerLike from zarr.core.chunk_key_encodings import DefaultChunkKeyEncoding @@ -281,96 +280,6 @@ def _sharding_codecs( ) -def _divisors(n: int) -> st.SearchStrategy[int]: - return st.sampled_from([d for d in range(1, n + 1) if n % d == 0]) - - -@st.composite -def transposed_sharding_chains( - draw: st.DrawFn, *, shape: tuple[int, ...] -) -> tuple[tuple[int, ...], list[Codec] | None, ShardingCodec, bool]: - """A codec chain with a `TransposeCodec` ahead of a `ShardingCodec`, and its validity. - - Returns `(chunks, filters, serializer, valid)` for - `zarr.create_array(shape=shape, chunks=chunks, filters=filters, serializer=serializer)`. - One of three layouts is drawn: - - - `transpose then shard`: `filters=[transpose]`, the serializer shards the - transposed chunk. - - `nested shard, transpose between levels`: an outer `ShardingCodec` whose - inner chain is `[transpose, ShardingCodec]`. - - `transpose inside shard`: a `ShardingCodec` whose inner chain is - `[transpose, BytesCodec]`; a control that is always valid. - - A sharding codec placed after the transpose sees transposed chunks, so its - `chunk_shape` must divide the transposed edges, not the array's chunk - grid. That inner chunk shape is drawn without regard to divisibility, so - about half of the drawn chains are invalid; `valid` is the oracle: every - edge of the sharding codec's chunk shape divides the transposed edge it - applies to. Any shard shape drawn for the outer codec always divides the - array's chunk grid, so `valid` is decided by the transpose alone. - """ - ndim = len(shape) - # Edges are drawn uniformly rather than via ``chunk_shapes`` (which favors - # many chunks of edge 1): unequal edges are what make a transpose change - # the chunk shape, which is the case this strategy exists for. - chunks = tuple(draw(st.integers(min_value=1, max_value=s)) for s in shape) - order = tuple(draw(st.permutations(range(ndim)), label="transpose order")) - transpose = TransposeCodec(order=order) - layout = draw( - st.sampled_from( - [ - "transpose then shard", - "nested shard, transpose between levels", - "transpose inside shard", - ] - ), - label="layout", - ) - event("transposed sharding chain layout", layout) - - def transposed(edges: tuple[int, ...]) -> tuple[int, ...]: - return tuple(edges[order[i]] for i in range(ndim)) - - event("transpose changes the chunk shape", "yes" if transposed(chunks) != chunks else "no") - - def draw_inner(seen: tuple[int, ...]) -> tuple[tuple[int, ...], bool]: - # Half the chains are valid (every edge a divisor of the transposed - # edge); the other half break exactly one axis with a non-divisor, of - # which `t + 1` guarantees at least one exists. - valid = draw(st.booleans(), label="valid chain") - broken = None if valid else draw(st.integers(min_value=0, max_value=ndim - 1)) - inner = tuple( - draw( - _divisors(t) - if axis != broken - else st.sampled_from([i for i in range(1, t + 2) if t % i != 0]) - ) - for axis, t in enumerate(seen) - ) - assert valid == all(t % i == 0 for t, i in zip(seen, inner, strict=True)) - return inner, valid - - filters: list[Codec] | None = None - if layout == "transpose then shard": - inner, valid = draw_inner(transposed(chunks)) - filters = [transpose] - serializer = ShardingCodec(chunk_shape=inner, codecs=[BytesCodec()]) - else: - mid = tuple(draw(_divisors(c)) for c in chunks) - if layout == "nested shard, transpose between levels": - inner, valid = draw_inner(transposed(mid)) - serializer = ShardingCodec( - chunk_shape=mid, - codecs=[transpose, ShardingCodec(chunk_shape=inner, codecs=[BytesCodec()])], - ) - else: - valid = True - serializer = ShardingCodec(chunk_shape=mid, codecs=[transpose, BytesCodec()]) - event("transposed sharding chain", "valid" if valid else "invalid") - return chunks, filters, serializer, valid - - @st.composite def np_array_and_chunks( draw: st.DrawFn, diff --git a/tests/test_codecs/test_codec_chain_validation.py b/tests/test_codecs/test_codec_chain_validation.py index 2327398c7e..eb799b1944 100644 --- a/tests/test_codecs/test_codec_chain_validation.py +++ b/tests/test_codecs/test_codec_chain_validation.py @@ -147,10 +147,8 @@ def test_sharding_inner_chain_is_validated() -> None: chunk_shape=CHUNKS, codecs=(ReshapeCodec(shape=(2, 3, 2, 2)), TransposeCodec(order=(2, 1, 0))), ) - with pytest.raises(ValueError, match="`order` tuple must have as many entries") as exc_info: + with pytest.raises(ValueError, match="`order` tuple must have as many entries"): _metadata((bad,), chunk_shape=SHAPE) - # the error names the codec that raised and the shape it was checked against - assert any("TransposeCodec" in n and "(2, 3, 2, 2)" in n for n in exc_info.value.__notes__) def _rectilinear_transpose_sharding_metadata(inner: tuple[int, int]) -> ArrayV3Metadata: diff --git a/tests/test_codecs/test_codec_chain_validation_properties.py b/tests/test_codecs/test_codec_chain_validation_properties.py index 4f839ed286..175a47d75e 100644 --- a/tests/test_codecs/test_codec_chain_validation_properties.py +++ b/tests/test_codecs/test_codec_chain_validation_properties.py @@ -1,22 +1,15 @@ """Property-based tests for codec-chain validation with shape-changing codecs. -Two invariants are tested against explicit oracles: - -1. Acceptance implies round-trip: any reshape+transpose chain that metadata - validation accepts must encode and decode data losslessly (and its metadata - must survive JSON serialization), while a transpose order of the wrong rank - must be rejected. - -2. For a rectilinear grid followed by a shape-changing codec and a - size-sensitive codec (sharding), acceptance must exactly equal the oracle - "every chunk shape in the grid, transformed by the chain, satisfies the - size constraint" — not just the largest chunk (see - ``evolve_and_validate_codecs``). +Acceptance implies round-trip: any reshape+transpose chain that metadata +validation accepts must encode and decode data losslessly (and its metadata +must survive JSON serialization), while a transpose order of the wrong rank +must be rejected. A fill-changing inner codec must be validated against the +actual fill value. (Transpose ahead of sharding, regular and nested, is +covered by `test_transposed_sharding_chain_validation` in `test_properties.py`.) """ from __future__ import annotations -import itertools import math from typing import TYPE_CHECKING @@ -27,9 +20,8 @@ from collections.abc import Iterator import zarr -from zarr.codecs import ShardingCodec, TransposeCodec -from zarr.core.dtype import Int32 -from zarr.core.metadata.v3 import ArrayV3Metadata, RectilinearChunkGridMetadata +from zarr.codecs import TransposeCodec +from zarr.core.metadata.v3 import ArrayV3Metadata from zarr.registry import _codec_registries, register_codec from .test_codec_chain_validation import ReshapeCodec @@ -129,63 +121,6 @@ def test_wrong_rank_transpose_after_reshape_rejected( ) -@st.composite -def rectilinear_transpose_sharding_cases( - draw: st.DrawFn, -) -> tuple[tuple[int | tuple[int, ...], ...], tuple[int, ...], tuple[int, ...]]: - """(rectilinear chunk_shapes, transpose order, inner shard shape).""" - ndim = draw(st.integers(min_value=2, max_value=3)) - chunk_shapes: list[int | tuple[int, ...]] = [] - for _ in range(ndim): - edges = draw(st.lists(st.integers(min_value=1, max_value=6), min_size=1, max_size=3)) - # exercise the bare-int (uniform edge) spelling as well - if len(edges) == 1 and draw(st.booleans()): - chunk_shapes.append(edges[0]) - else: - chunk_shapes.append(tuple(edges)) - order = tuple(draw(st.permutations(range(ndim)))) - inner = tuple(draw(st.integers(min_value=1, max_value=6)) for _ in range(ndim)) - return tuple(chunk_shapes), order, inner - - -@settings(deadline=None) -@given(case=rectilinear_transpose_sharding_cases()) -def test_rectilinear_transpose_sharding_matches_oracle( - case: tuple[tuple[int | tuple[int, ...], ...], tuple[int, ...], tuple[int, ...]], -) -> None: - """transpose-then-shard over a rectilinear grid is accepted exactly when - every transposed chunk shape is divisible by the inner shard shape.""" - chunk_shapes, order, inner = case - per_dim = tuple((e,) if isinstance(e, int) else e for e in chunk_shapes) - oracle_ok = all( - all(chunk[order[i]] % inner[i] == 0 for i in range(len(inner))) - for chunk in itertools.product(*per_dim) - ) - # array shape: bare-int (uniform) edges cover any extent; explicit edge - # lists must sum to at least the extent. - shape = tuple(e if isinstance(e, int) else sum(e) for e in chunk_shapes) - - def build() -> ArrayV3Metadata: - return ArrayV3Metadata( - shape=shape, - data_type=Int32(), - chunk_grid=RectilinearChunkGridMetadata(chunk_shapes=chunk_shapes), - chunk_key_encoding={"name": "default"}, - fill_value=0, - codecs=(TransposeCodec(order=order), ShardingCodec(chunk_shape=inner)), - attributes=None, - dimension_names=None, - ) - - with zarr.config.set({"array.rectilinear_chunks": True}): - if oracle_ok: - meta = build() - assert ArrayV3Metadata.from_dict(meta.to_dict()) == meta - else: - with pytest.raises(ValueError, match="not\\s+divisible"): - build() - - @given(offset=st.integers(min_value=1, max_value=254), sharded=st.booleans()) def test_inner_validation_uses_the_actual_fill_value(offset: int, sharded: bool) -> None: """Shard validation must not resolve a fill-changing codec against a made-up zero.""" diff --git a/tests/test_properties.py b/tests/test_properties.py index 5c3f480f46..2091fa251f 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -17,7 +17,9 @@ import hypothesis.strategies as st from hypothesis import assume, event, given, settings +from zarr.abc.codec import Codec from zarr.abc.store import Store +from zarr.codecs import BytesCodec, ShardingCodec, TransposeCodec from zarr.core.common import ZARR_JSON, ZARRAY_JSON, ZATTRS_JSON from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata from zarr.core.sync import sync @@ -34,7 +36,6 @@ sharded_arrays, simple_arrays, stores, - transposed_sharding_chains, zarr_formats, ) @@ -499,6 +500,96 @@ def test_chunks_param_from_rectilinear_bare_int_roundtrip() -> None: assert dst.metadata.chunk_grid == grid # type: ignore[union-attr] +def _divisors(n: int) -> st.SearchStrategy[int]: + return st.sampled_from([d for d in range(1, n + 1) if n % d == 0]) + + +@st.composite +def transposed_sharding_chains( + draw: st.DrawFn, *, shape: tuple[int, ...] +) -> tuple[tuple[int, ...], list[Codec] | None, ShardingCodec, bool]: + """A codec chain with a `TransposeCodec` ahead of a `ShardingCodec`, and its validity. + + Returns `(chunks, filters, serializer, valid)` for + `zarr.create_array(shape=shape, chunks=chunks, filters=filters, serializer=serializer)`. + One of three layouts is drawn: + + - `transpose then shard`: `filters=[transpose]`, the serializer shards the + transposed chunk. + - `nested shard, transpose between levels`: an outer `ShardingCodec` whose + inner chain is `[transpose, ShardingCodec]`. + - `transpose inside shard`: a `ShardingCodec` whose inner chain is + `[transpose, BytesCodec]`; a control that is always valid. + + A sharding codec placed after the transpose sees transposed chunks, so its + `chunk_shape` must divide the transposed edges, not the array's chunk + grid. That inner chunk shape is drawn without regard to divisibility, so + about half of the drawn chains are invalid; `valid` is the oracle: every + edge of the sharding codec's chunk shape divides the transposed edge it + applies to. Any shard shape drawn for the outer codec always divides the + array's chunk grid, so `valid` is decided by the transpose alone. + """ + ndim = len(shape) + # Edges are drawn uniformly rather than via ``chunk_shapes`` (which favors + # many chunks of edge 1): unequal edges are what make a transpose change + # the chunk shape, which is the case this strategy exists for. + chunks = tuple(draw(st.integers(min_value=1, max_value=s)) for s in shape) + order = tuple(draw(st.permutations(range(ndim)), label="transpose order")) + transpose = TransposeCodec(order=order) + layout = draw( + st.sampled_from( + [ + "transpose then shard", + "nested shard, transpose between levels", + "transpose inside shard", + ] + ), + label="layout", + ) + event("transposed sharding chain layout", layout) + + def transposed(edges: tuple[int, ...]) -> tuple[int, ...]: + return tuple(edges[order[i]] for i in range(ndim)) + + event("transpose changes the chunk shape", "yes" if transposed(chunks) != chunks else "no") + + def draw_inner(seen: tuple[int, ...]) -> tuple[tuple[int, ...], bool]: + # Half the chains are valid (every edge a divisor of the transposed + # edge); the other half break exactly one axis with a non-divisor, of + # which `t + 1` guarantees at least one exists. + valid = draw(st.booleans(), label="valid chain") + broken = None if valid else draw(st.integers(min_value=0, max_value=ndim - 1)) + inner = tuple( + draw( + _divisors(t) + if axis != broken + else st.sampled_from([i for i in range(1, t + 2) if t % i != 0]) + ) + for axis, t in enumerate(seen) + ) + assert valid == all(t % i == 0 for t, i in zip(seen, inner, strict=True)) + return inner, valid + + filters: list[Codec] | None = None + if layout == "transpose then shard": + inner, valid = draw_inner(transposed(chunks)) + filters = [transpose] + serializer = ShardingCodec(chunk_shape=inner, codecs=[BytesCodec()]) + else: + mid = tuple(draw(_divisors(c)) for c in chunks) + if layout == "nested shard, transpose between levels": + inner, valid = draw_inner(transposed(mid)) + serializer = ShardingCodec( + chunk_shape=mid, + codecs=[transpose, ShardingCodec(chunk_shape=inner, codecs=[BytesCodec()])], + ) + else: + valid = True + serializer = ShardingCodec(chunk_shape=mid, codecs=[transpose, BytesCodec()]) + event("transposed sharding chain", "valid" if valid else "invalid") + return chunks, filters, serializer, valid + + @given(data=st.data()) @pytest.mark.filterwarnings( "ignore:Combining a `sharding_indexed` codec:zarr.errors.ZarrUserWarning" From aaefe61cb915e043dc9147ddf174e64835b86e6f Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 14 Sep 2026 12:30:41 +0200 Subject: [PATCH 08/11] chore: renumber changelog fragment to upstream PR 4352 Assisted-by: ClaudeCode:claude-fable-5-1 Co-Authored-By: Claude Fable 5.1 --- changes/{303.bugfix.md => 4352.bugfix.md} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename changes/{303.bugfix.md => 4352.bugfix.md} (100%) diff --git a/changes/303.bugfix.md b/changes/4352.bugfix.md similarity index 100% rename from changes/303.bugfix.md rename to changes/4352.bugfix.md From b87fd36ce272f73935d56be049c9056d6375206e Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 14 Sep 2026 12:57:58 +0200 Subject: [PATCH 09/11] fix(metadata): validate factored grids without chunk enumeration Preserve grid geometry for identity, cast, scale, and transpose codecs. Stream arbitrary resolver chains without retaining the cross product, and test bounded work for common paths. Assisted-by: Codex:GPT-6 --- changes/4352.bugfix.md | 2 +- src/zarr/core/metadata/v3.py | 131 +++++++++++------- .../test_codec_chain_validation.py | 72 ++++++++++ 3 files changed, 153 insertions(+), 52 deletions(-) diff --git a/changes/4352.bugfix.md b/changes/4352.bugfix.md index 01304976c0..21c003bee0 100644 --- a/changes/4352.bugfix.md +++ b/changes/4352.bugfix.md @@ -1 +1 @@ -Codec chains are now validated against the chunk spec threaded through each codec's `resolve_metadata`, the same way the codec pipeline resolves it at encode time, instead of against the array-level shape. Previously, an `array -> array` codec that changes a chunk's shape or rank (such as the zarr-extensions `reshape` codec) followed by a codec whose configuration refers to the transformed chunk (e.g. `transpose`) was wrongly rejected. The sharding codec now validates its inner codec chain in the same way. Inner sharding chains are validated during evolution using the actual fill value, so valid fill-changing codecs are not rejected against a fabricated default. \ No newline at end of file +Validate codec chains against evolving chunk specs, including inner sharding chains and the actual fill value. Preserve factored grid geometry for default metadata resolvers, cast/scale codecs, and transpose so common chains validate per-axis sizes without enumerating chunk combinations. Stream arbitrary shape-changing chains without materializing the cross product. diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index a4767e309c..17c29eedcb 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -2,13 +2,13 @@ import itertools import json -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass, field, replace from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, TypeGuard, cast from typing_extensions import TypedDict -from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec, Codec +from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BaseCodec, BytesBytesCodec, Codec from zarr.abc.metadata import Metadata from zarr.core._json import json_to_buffer from zarr.core.array_spec import ArrayConfig, ArraySpec @@ -132,15 +132,15 @@ def representative_chunk_shape(chunk_grid: ChunkGridMetadata) -> tuple[int, ...] Regular grids have exactly one chunk shape. Rectilinear grids have many; the largest edge along each dimension is used. This is only suitable where a single shape is structurally required (rank checks, codec evolution); - size-sensitive validation must consider every distinct chunk shape, see - `_distinct_chunk_shapes`. + size-sensitive validation must cover every chunk size, either through the + factored grid or by streaming `_distinct_chunk_shapes`. """ if isinstance(chunk_grid, RegularChunkGridMetadata): return chunk_grid.chunk_shape return tuple(s if isinstance(s, int) else max(s) for s in chunk_grid.chunk_shapes) -def _distinct_chunk_shapes(chunk_grid: ChunkGridMetadata) -> list[tuple[int, ...]]: +def _distinct_chunk_shapes(chunk_grid: ChunkGridMetadata) -> Iterator[tuple[int, ...]]: """Every distinct chunk shape occurring in `chunk_grid`. For a rectilinear grid every combination of per-dimension distinct edges @@ -148,11 +148,23 @@ def _distinct_chunk_shapes(chunk_grid: ChunkGridMetadata) -> list[tuple[int, ... edge along every other), so this is the full cross product. """ if isinstance(chunk_grid, RegularChunkGridMetadata): - return [chunk_grid.chunk_shape] + yield chunk_grid.chunk_shape + return per_dim = ( (s,) if isinstance(s, int) else tuple(dict.fromkeys(s)) for s in chunk_grid.chunk_shapes ) - return list(itertools.product(*per_dim)) + yield from itertools.product(*per_dim) + + +def _resolved_chunk_specs( + chunk_grid: ChunkGridMetadata, spec: ArraySpec, codecs: tuple[Codec, ...] +) -> Iterator[ArraySpec]: + """Stream chunk specs through a codec prefix without storing the cross product.""" + for shape in _distinct_chunk_shapes(chunk_grid): + current = replace(spec, shape=shape) + for codec in codecs: + current = codec.resolve_metadata(current) + yield current def evolve_and_validate_codecs( @@ -162,59 +174,76 @@ def evolve_and_validate_codecs( chunk_grid: ChunkGridMetadata, chunk_spec: ArraySpec, ) -> tuple[Codec, ...]: - """Evolve and validate a codec chain, threading the chunk spec. - - Each codec is evolved and validated against the chunk spec produced by the - previous codec's `resolve_metadata`, the same spec it will see at - encode/decode time, not against the array-level metadata. Earlier - array->array codecs may change the dtype (`cast_value`) or the shape and - even the rank of a chunk (the `reshape` extension codec, which the spec - explicitly allows to be followed by `transpose`). - - `shape` and `chunk_grid` are the array-level values passed to - `Codec.validate`. They are handed unchanged to every codec until one - changes the shape of any chunk; from then on the array-level values are no - longer meaningful for the remaining codecs. Because `validate` checks may - be size-sensitive (sharding divisibility), every *distinct* chunk shape of - the grid is threaded through `resolve_metadata` and validated - individually; for a rectilinear grid, a single representative shape would - not be sound: an inner chunk size that divides the largest chunk need not - divide the others. Each threaded shape is presented to `validate` as a - regular grid of that shape, the only shape-related facts that survive a - per-chunk transformation. - - `chunk_spec` (built from the representative chunk shape) is threaded - separately as the single spec used for codec evolution and dtype tracking, - since evolution must produce one codec chain. - - Per-codec `validate` runs before `resolve_metadata`, since the latter - may rely on invariants the former checks (e.g. `cast_value` rejects - complex source dtypes that would otherwise crash `_do_cast`). + """Evolve and validate codecs while preserving factored chunk geometry. + + Codecs with the default identity metadata resolver retain the grid. + CastValue and ScaleOffset also retain its geometry; TransposeCodec permutes + the dimensions of both the shape and grid. These paths resolve one spec per + codec and validate the grid directly, without enumerating chunk combinations. + Only the exact built-in classes use specialized geometry rules; overridden + metadata resolvers must use the general path. + + Arbitrary metadata resolvers stream each distinct chunk spec. Validation + uses the array geometry until a chunk shape changes, then each resolved + shape is validated as a regular grid. Full specs retain per-chunk dtype and + fill values. Prefixes are replayed for subsequent codecs rather than cached: + this bounds memory but may still require combinatorial work for arbitrary + transforms. Evolution uses one representative spec to produce one chain. + + Validation precedes metadata resolution, whose implementation may require + the validated dtype and geometry. Each codec is evolved only once here. """ + from zarr.codecs.cast_value import CastValue + from zarr.codecs.scale_offset import ScaleOffset + from zarr.codecs.transpose import TransposeCodec + out: list[Codec] = [] spec = chunk_spec - threaded = _distinct_chunk_shapes(chunk_grid) + grid = chunk_grid + validation_shape = shape + fallback_grid: ChunkGridMetadata | None = None + fallback_spec = chunk_spec + fallback_codecs: tuple[Codec, ...] = () shapes_changed = False for codec in codecs: evolved = codec.evolve_from_array_spec(spec) - # The array-level shape and grid are handed to validate until a codec - # changes chunk shapes; from then on each distinct chunk shape is - # validated as its own regular grid. - if shapes_changed: - for s in threaded: + preserves_shape = type(evolved).resolve_metadata is BaseCodec.resolve_metadata or type( + evolved + ) in (CastValue, ScaleOffset) + if fallback_grid is None and (preserves_shape or type(evolved) is TransposeCodec): + evolved.validate(shape=validation_shape, dtype=spec.dtype, chunk_grid=grid) + spec = evolved.resolve_metadata(spec) + if type(evolved) is TransposeCodec: + order = evolved.order + validation_shape = tuple(validation_shape[d] for d in order) + if isinstance(grid, RegularChunkGridMetadata): + grid = RegularChunkGridMetadata( + chunk_shape=tuple(grid.chunk_shape[d] for d in order) + ) + else: + grid = RectilinearChunkGridMetadata( + chunk_shapes=tuple(grid.chunk_shapes[d] for d in order) + ) + else: + if fallback_grid is None: + fallback_grid, fallback_spec = grid, spec + changed = False + for current in _resolved_chunk_specs(fallback_grid, fallback_spec, fallback_codecs): evolved.validate( - shape=s, dtype=spec.dtype, chunk_grid=RegularChunkGridMetadata(chunk_shape=s) + shape=current.shape if shapes_changed else validation_shape, + dtype=current.dtype, + chunk_grid=( + RegularChunkGridMetadata(chunk_shape=current.shape) + if shapes_changed + else grid + ), ) - else: - evolved.validate(shape=shape, dtype=spec.dtype, chunk_grid=chunk_grid) + resolved = evolved.resolve_metadata(current) + changed |= resolved.shape != current.shape + shapes_changed |= changed + fallback_codecs += (evolved,) + spec = evolved.resolve_metadata(spec) out.append(evolved) - resolved = list( - dict.fromkeys(evolved.resolve_metadata(replace(spec, shape=s)).shape for s in threaded) - ) - if resolved != threaded: - shapes_changed = True - threaded = resolved - spec = evolved.resolve_metadata(spec) return tuple(out) diff --git a/tests/test_codecs/test_codec_chain_validation.py b/tests/test_codecs/test_codec_chain_validation.py index eb799b1944..aa628ad607 100644 --- a/tests/test_codecs/test_codec_chain_validation.py +++ b/tests/test_codecs/test_codec_chain_validation.py @@ -186,3 +186,75 @@ def test_rectilinear_every_chunk_shape_validated() -> None: pytest.raises(ValueError, match="not\\s+divisible"), ): _rectilinear_transpose_sharding_metadata((5, 3)) + + +@pytest.mark.parametrize("rank", [4, 12]) +@pytest.mark.parametrize("chain", ["bytes", "scale-offset", "transpose-sharding"]) +def test_factored_validation_does_not_visit_chunk_combinations( + rank: int, chain: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """Common codec chains resolve specs independently of the chunk cross product.""" + from zarr.abc.codec import BaseCodec + from zarr.codecs.scale_offset import ScaleOffset + + calls = 0 + original = BaseCodec.resolve_metadata + + def counted(self: Any, spec: ArraySpec) -> ArraySpec: + nonlocal calls + calls += 1 + assert calls <= 10, "metadata validation enumerated chunk combinations" + return original(self, spec) + + monkeypatch.setattr(BaseCodec, "resolve_metadata", counted) + codecs: tuple[Any, ...] + if chain == "bytes": + codecs = (BytesCodec(),) + elif chain == "scale-offset": + codecs = (ScaleOffset(offset=0), BytesCodec()) + else: + codecs = ( + TransposeCodec(order=tuple(reversed(range(rank)))), + ShardingCodec(chunk_shape=(1,) * rank), + ) + with zarr.config.set({"array.rectilinear_chunks": True}): + ArrayV3Metadata( + shape=(3,) * rank, + data_type=Int32(), + chunk_grid=RectilinearChunkGridMetadata(chunk_shapes=((1, 2),) * rank), + chunk_key_encoding={"name": "default"}, + fill_value=0, + codecs=codecs, + attributes={}, + dimension_names=None, + ) + assert calls <= 10 + + +def test_arbitrary_shape_validation_stops_at_first_invalid_chunk( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A failing arbitrary transform does not consume all chunk combinations.""" + import itertools + + original_product = itertools.product + + def guarded_product(*args: Any) -> Iterator[tuple[int, ...]]: + yield next(original_product(*args)) + raise AssertionError("consumed combinations after the first invalid chunk") + + monkeypatch.setattr(itertools, "product", guarded_product) + with ( + zarr.config.set({"array.rectilinear_chunks": True}), + pytest.raises(ValueError, match="cannot reshape a chunk"), + ): + ArrayV3Metadata( + shape=(3, 3), + data_type=Int32(), + chunk_grid=RectilinearChunkGridMetadata(chunk_shapes=((1, 2), (1, 2))), + chunk_key_encoding={"name": "default"}, + fill_value=0, + codecs=(ReshapeCodec(shape=(4,)), BytesCodec()), + attributes={}, + dimension_names=None, + ) From 8a8f3b4ef1524837d59cf879010deab719c89a51 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 14 Sep 2026 17:30:06 +0200 Subject: [PATCH 10/11] fix(metadata): thread chunk grids through codec validation without enumerating chunks Validating codec chains on rectilinear grids streamed every combination of per-axis chunk edges through each codec whose `resolve_metadata` was overridden, e.g. numcodecs `delta`. A 3-d grid with 100 distinct edges per axis took ~3.8 s per metadata construction, growing as n**ndim. Follow zarrs: codecs map a whole chunk grid via the new `BaseCodec.resolve_chunk_grid`. Dtype/fill-value codecs declare the identity and transpose declares a permutation, so those chains stay exact. An undeclared codec is exact on regular grids; on rectilinear grids it makes the rest of the chain chunk-local, validated against one representative chunk, with other chunk shapes checked at encode/decode time. ShardingCodec now enforces divisibility at run time, which would otherwise floor-divide and silently corrupt data. Document the trade-off in the codec docstrings, the extending guide, the rectilinear sharding docs and the changelog. Assisted-by: ClaudeCode:claude-opus-5 --- changes/4352.bugfix.md | 6 +- docs/user-guide/arrays.md | 8 + docs/user-guide/extending.md | 56 ++++ src/zarr/abc/codec.py | 49 +++ src/zarr/codecs/cast_value.py | 6 + src/zarr/codecs/numcodecs/_codecs.py | 19 ++ src/zarr/codecs/scale_offset.py | 6 + src/zarr/codecs/sharding.py | 28 +- src/zarr/codecs/transpose.py | 15 + src/zarr/core/metadata/v3.py | 193 ++++++------ .../test_codec_chain_validation.py | 278 +++++++++++++----- 11 files changed, 500 insertions(+), 164 deletions(-) diff --git a/changes/4352.bugfix.md b/changes/4352.bugfix.md index 21c003bee0..c52e343521 100644 --- a/changes/4352.bugfix.md +++ b/changes/4352.bugfix.md @@ -1 +1,5 @@ -Validate codec chains against evolving chunk specs, including inner sharding chains and the actual fill value. Preserve factored grid geometry for default metadata resolvers, cast/scale codecs, and transpose so common chains validate per-axis sizes without enumerating chunk combinations. Stream arbitrary shape-changing chains without materializing the cross product. +Validate each codec in a chain against the chunk geometry produced by the codecs before it, including the inner codecs of a sharding codec and the actual fill value. For example, a sharding codec placed after a transpose must now divide the transposed chunks; previously it was checked against the untransposed chunk grid, which could accept chains that read back wrong data and reject valid ones. + +Geometry is threaded through the chain as a whole chunk grid, so validation cost no longer depends on the number of distinct chunk shapes. Codecs describe how they map the grid with the new `BaseCodec.resolve_chunk_grid` method; the built-in dtype and fill-value codecs (`cast_value`, `scale_offset`, and the numcodecs `delta`, `fixedscaleoffset` and `astype` filters) declare the identity and `transpose` declares a permutation. + +Trade-off: on a rectilinear chunk grid, a codec that overrides `resolve_metadata` without implementing `resolve_chunk_grid` makes the rest of its chain "chunk-local". Codecs after it are validated against one representative chunk (the largest edge along each axis), so a chain that is invalid only for some other chunk shape is accepted when the array is created and fails when such a chunk is first written or read. `ShardingCodec` now checks divisibility at encode and decode time, so this case raises an error instead of corrupting data. Regular chunk grids are always validated exactly. See "Chunk geometry and validation" in the extending guide. diff --git a/docs/user-guide/arrays.md b/docs/user-guide/arrays.md index 140bc4bd09..1532ebbdb0 100644 --- a/docs/user-guide/arrays.md +++ b/docs/user-guide/arrays.md @@ -749,6 +749,14 @@ print(z[50:70, 40:60]) Note that rectilinear inner chunks with sharding are not supported — only the shard boundaries can be rectilinear. +!!! note "When the divisibility check happens" + Zarr checks that the inner chunk shape divides every shard when the array is + created, as long as each filter before the sharding codec describes how it maps + the chunk grid. All built-in filters that keep the chunk shape, and `transpose`, do. A third-party filter that changes chunk + metadata without describing its grid limits that check to the largest shard. A + smaller shard that is not divisible is then reported only when it is first written + or read. See [Chunk geometry and validation](extending.md#chunk-geometry-and-validation). + For such arrays, `.chunks` returns the (regular) inner chunk shape, while `.shards` raises `NotImplementedError` since there is no single uniform shard shape — use `.write_chunk_sizes` for the per-dimension shard sizes. `.info` diff --git a/docs/user-guide/extending.md b/docs/user-guide/extending.md index 507afedea7..68ca4602af 100644 --- a/docs/user-guide/extending.md +++ b/docs/user-guide/extending.md @@ -45,9 +45,65 @@ Custom codecs should also implement the following methods: array metadata. It should raise errors if not. - `resolve_metadata` (optional), which is important for codecs that change the shape, dtype or fill value of a chunk. +- `resolve_chunk_grid` (optional, but recommended for any codec that overrides + `resolve_metadata`), which describes how the codec changes the chunk grid as a whole. + See [Chunk geometry and validation](#chunk-geometry-and-validation). - `evolve_from_array_spec` (optional), which can be useful for automatically filling in codec configuration metadata from the array metadata. +### Chunk geometry and validation + +When array metadata is created, Zarr validates every codec against the chunks it will +actually receive, which are the chunks produced by the codecs before it. For example, in +the chain `transpose` → `sharding_indexed`, the shard's inner chunk shape must divide the +*transposed* chunks. + +The geometry is carried through the chain as a whole chunk grid rather than chunk by +chunk. A rectilinear grid with `n` distinct edge lengths along each of `d` axes has `n**d` +distinct chunk shapes, which quickly becomes too many to check one at a time (a 3-d grid +with 1000 distinct edges per axis has a billion). Each codec therefore reports how it +changes the grid through `resolve_chunk_grid(shape=..., chunk_grid=...)`: + +- **Return `(shape, chunk_grid)` unchanged** if the codec never changes the chunk + shape, even if it changes the data type or fill value. This is the default for codecs + that do not override `resolve_metadata`. If your codec overrides `resolve_metadata` + only to change the data type or fill value, override `resolve_chunk_grid` too, as + `CastValue`, `ScaleOffset` and the numcodecs `Delta` do. +- **Return the mapped shape and grid** if the codec changes chunk shapes in a way a grid + can express. `TransposeCodec`, for example, permutes the axes of both. +- **Return `None`** if the chunks after the codec cannot be described by a single grid + computed from the input grid. This is the default for codecs that override + `resolve_metadata`. + +Returning `None` has no cost on a regular chunk grid, because every chunk has the same +shape and `resolve_metadata` describes it exactly. On a rectilinear chunk grid it makes +the rest of the chain **chunk-local**. From then on, codecs are validated against a single +representative chunk: the one with the largest edge along each axis. This has the +following consequences: + +- **Rejections are always correct.** The representative is a real chunk of the array, so + a codec that rejects it would fail on that chunk. +- **Acceptance is incomplete.** A chain that is invalid only for some *other* chunk + shape is accepted when the array is created, and the error surfaces when a chunk of + that shape is first encoded or decoded. Suppose a sharding codec follows an undeclared + filter on a grid with chunk edges `[4, 6]`. An inner chunk size of `3` divides the + representative edge `6`, so the array is created, but writing to a chunk with edge `4` + raises an error. +- **Size-sensitive codecs must check at run time.** Any codec whose correctness depends on + the chunk shape (for example, requiring it to divide evenly) must repeat that check + when encoding and decoding, because metadata validation may have seen only the + representative chunk. `ShardingCodec` does this. A codec that skips the check can + silently corrupt data on the chunks that validation never saw. + +Zarr uses a declared grid only when it can trust it. A declaration is ignored, and the +codec treated as returning `None`, if a subclass overrides `resolve_metadata` without +also overriding `resolve_chunk_grid`. It is also ignored if the declared grid disagrees +with what `resolve_metadata` returns for the representative chunk. + +This design follows [zarrs](https://github.com/zarrs/zarrs), whose array-to-array codecs +map chunk grids through `encoded_chunk_grid` and return a "chunk-local" result when no +whole-array grid exists. + To use custom codecs in Zarr, they need to be registered using the [entrypoint mechanism](https://packaging.python.org/en/latest/specifications/entry-points/). Commonly, entrypoints are declared in the `pyproject.toml` of your package under the diff --git a/src/zarr/abc/codec.py b/src/zarr/abc/codec.py index 34d349e6d1..e567a8b482 100644 --- a/src/zarr/abc/codec.py +++ b/src/zarr/abc/codec.py @@ -148,6 +148,55 @@ def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec: """ return chunk_spec + def resolve_chunk_grid( + self, *, shape: tuple[int, ...], chunk_grid: ChunkGridMetadata + ) -> tuple[tuple[int, ...], ChunkGridMetadata] | None: + """The array shape and chunk grid seen by the codecs after this one. + + This is the whole-array counterpart of `resolve_metadata`: where + `resolve_metadata` maps the spec of one chunk, this maps the geometry + of every chunk at once. It is used to validate a codec chain when the + array metadata is created, so that a later size-sensitive codec (such + as `ShardingCodec`) is checked against the chunks it will actually + receive, without enumerating the chunks of the grid. + + Return `None` when the chunks after this codec cannot be described by + a single chunk grid computed from `chunk_grid` alone. On a regular grid + that costs nothing, because all chunks have one shape and + `resolve_metadata` describes them exactly. On a rectilinear grid it + makes the remaining chain "chunk-local": later codecs are validated + against one representative chunk only, so a chain that is invalid for + some other chunk shape is accepted at creation and fails when such a + chunk is first encoded or decoded (see + `zarr.core.metadata.v3.evolve_and_validate_codecs`). + + The default declares the identity when `resolve_metadata` is not + overridden, and returns `None` otherwise. Codecs that override + `resolve_metadata` without changing the chunk shape (for example to + change the data type or fill value) should override this method to + return `(shape, chunk_grid)` unchanged, and codecs that change the chunk + shape in a way expressible as a grid (for example a permutation of the + axes) should return the mapped geometry. + + A declaration is ignored if a subclass overrides `resolve_metadata` + without also overriding this method, or if it disagrees with + `resolve_metadata` on the representative chunk. + + Parameters + ---------- + shape : tuple[int, ...] + The array shape seen by this codec. + chunk_grid : ChunkGridMetadata + The chunk grid seen by this codec. + + Returns + ------- + tuple[tuple[int, ...], ChunkGridMetadata] | None + """ + if type(self).resolve_metadata is BaseCodec.resolve_metadata: + return shape, chunk_grid + return None + def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self: """Fills in codec configuration parameters that can be automatically inferred from the array metadata. diff --git a/src/zarr/codecs/cast_value.py b/src/zarr/codecs/cast_value.py index b19a10c873..2359786041 100644 --- a/src/zarr/codecs/cast_value.py +++ b/src/zarr/codecs/cast_value.py @@ -377,6 +377,12 @@ def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec: return replace(chunk_spec, dtype=target_zdtype, fill_value=new_fill) + def resolve_chunk_grid( + self, *, shape: tuple[int, ...], chunk_grid: ChunkGridMetadata + ) -> tuple[tuple[int, ...], ChunkGridMetadata]: + """Casting changes the data type and fill value, never the chunk shape.""" + return shape, chunk_grid + def _encode_sync( self, chunk_array: NDBuffer, diff --git a/src/zarr/codecs/numcodecs/_codecs.py b/src/zarr/codecs/numcodecs/_codecs.py index f44c35964c..cf32da12f2 100644 --- a/src/zarr/codecs/numcodecs/_codecs.py +++ b/src/zarr/codecs/numcodecs/_codecs.py @@ -48,6 +48,7 @@ from zarr.abc.numcodec import Numcodec from zarr.core.array_spec import ArraySpec from zarr.core.buffer import Buffer, NDBuffer + from zarr.core.metadata.v3 import ChunkGridMetadata CODEC_PREFIX = "numcodecs." @@ -241,6 +242,12 @@ def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec: return replace(chunk_spec, dtype=dtype) return chunk_spec + def resolve_chunk_grid( + self, *, shape: tuple[int, ...], chunk_grid: ChunkGridMetadata + ) -> tuple[tuple[int, ...], ChunkGridMetadata]: + """Delta encoding may change the data type, never the chunk shape.""" + return shape, chunk_grid + class BitRound(_NumcodecsArrayArrayCodec, codec_name="bitround"): pass @@ -253,6 +260,12 @@ def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec: return replace(chunk_spec, dtype=dtype) return chunk_spec + def resolve_chunk_grid( + self, *, shape: tuple[int, ...], chunk_grid: ChunkGridMetadata + ) -> tuple[tuple[int, ...], ChunkGridMetadata]: + """Scaling may change the data type, never the chunk shape.""" + return shape, chunk_grid + def evolve_from_array_spec(self, array_spec: ArraySpec) -> FixedScaleOffset: if self.codec_config.get("dtype") is None: dtype = array_spec.dtype.to_native_dtype() @@ -293,6 +306,12 @@ def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec: dtype = parse_dtype(np.dtype(self.codec_config["encode_dtype"]), zarr_format=3) # type: ignore[arg-type] return replace(chunk_spec, dtype=dtype) + def resolve_chunk_grid( + self, *, shape: tuple[int, ...], chunk_grid: ChunkGridMetadata + ) -> tuple[tuple[int, ...], ChunkGridMetadata]: + """Casting changes the data type, never the chunk shape.""" + return shape, chunk_grid + def evolve_from_array_spec(self, array_spec: ArraySpec) -> AsType: if self.codec_config.get("decode_dtype") is None: # TODO: remove these coverage exemptions the correct way, i.e. with tests diff --git a/src/zarr/codecs/scale_offset.py b/src/zarr/codecs/scale_offset.py index f2908da1b6..670af6b805 100644 --- a/src/zarr/codecs/scale_offset.py +++ b/src/zarr/codecs/scale_offset.py @@ -376,6 +376,12 @@ def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec: new_fill = _encode(fill, offset, scale) return replace(chunk_spec, fill_value=new_fill.reshape(()).item()) + def resolve_chunk_grid( + self, *, shape: tuple[int, ...], chunk_grid: ChunkGridMetadata + ) -> tuple[tuple[int, ...], ChunkGridMetadata]: + """Scaling changes the fill value, never the chunk shape.""" + return shape, chunk_grid + def _decode_sync( self, chunk_array: NDBuffer, diff --git a/src/zarr/codecs/sharding.py b/src/zarr/codecs/sharding.py index 46854e63e3..40eef39901 100644 --- a/src/zarr/codecs/sharding.py +++ b/src/zarr/codecs/sharding.py @@ -1597,14 +1597,28 @@ def _get_chunk_spec(self, shard_spec: ArraySpec) -> ArraySpec: ) def _get_chunks_per_shard(self, shard_spec: ArraySpec) -> tuple[int, ...]: - return tuple( - s // c - for s, c in zip( - shard_spec.shape, - self.chunk_shape, - strict=False, + """The number of inner chunks along each axis of a shard. + + Every encode, decode and size computation goes through here, so this is + also the run-time divisibility check. Metadata validation cannot always + establish divisibility for every shard: after a codec that maps chunk + shapes without declaring a chunk grid on a rectilinear grid, only one + representative shard is checked (see + `zarr.core.metadata.v3.evolve_and_validate_codecs`). Without this check + a non-dividing shard would be floor-divided and read or written with the + wrong layout, silently corrupting data. + """ + if len(shard_spec.shape) != len(self.chunk_shape) or any( + s % c != 0 for s, c in zip(shard_spec.shape, self.chunk_shape, strict=True) + ): + raise ValueError( + f"A shard of shape {shard_spec.shape} is not divisible by the shard's inner " + f"chunk shape {self.chunk_shape}. Metadata validation checks this for " + "every shard only when each codec before the sharding codec declares its " + "chunk grid (`resolve_chunk_grid`); otherwise it is detected here, when " + "such a shard is first encoded or decoded." ) - ) + return tuple(s // c for s, c in zip(shard_spec.shape, self.chunk_shape, strict=True)) def _shard_index_byte_range( self, chunks_per_shard: tuple[int, ...] diff --git a/src/zarr/codecs/transpose.py b/src/zarr/codecs/transpose.py index 5756fba2b4..d2458aa790 100644 --- a/src/zarr/codecs/transpose.py +++ b/src/zarr/codecs/transpose.py @@ -95,6 +95,21 @@ def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec: prototype=chunk_spec.prototype, ) + def resolve_chunk_grid( + self, *, shape: tuple[int, ...], chunk_grid: ChunkGridMetadata + ) -> tuple[tuple[int, ...], ChunkGridMetadata]: + """Permute the array shape and the per-axis chunk edges by `order`.""" + from zarr.core.metadata.v3 import RectilinearChunkGridMetadata, RegularChunkGridMetadata + + permuted_shape = tuple(shape[d] for d in self.order) + if isinstance(chunk_grid, RegularChunkGridMetadata): + return permuted_shape, RegularChunkGridMetadata( + chunk_shape=tuple(chunk_grid.chunk_shape[d] for d in self.order) + ) + return permuted_shape, RectilinearChunkGridMetadata( + chunk_shapes=tuple(chunk_grid.chunk_shapes[d] for d in self.order) + ) + def _decode_sync( self, chunk_array: NDBuffer, diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index 17c29eedcb..542fabcd11 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -1,14 +1,13 @@ from __future__ import annotations -import itertools import json -from collections.abc import Iterable, Iterator, Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass, field, replace from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, TypeGuard, cast from typing_extensions import TypedDict -from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BaseCodec, BytesBytesCodec, Codec +from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec, Codec from zarr.abc.metadata import Metadata from zarr.core._json import json_to_buffer from zarr.core.array_spec import ArrayConfig, ArraySpec @@ -130,41 +129,48 @@ def representative_chunk_shape(chunk_grid: ChunkGridMetadata) -> tuple[int, ...] """A single chunk shape standing in for every chunk of `chunk_grid`. Regular grids have exactly one chunk shape. Rectilinear grids have many; - the largest edge along each dimension is used. This is only suitable where - a single shape is structurally required (rank checks, codec evolution); - size-sensitive validation must cover every chunk size, either through the - factored grid or by streaming `_distinct_chunk_shapes`. + the largest edge along each dimension is used. Every combination of + per-dimension edges occurs in a rectilinear grid, so this is the shape of + a real chunk: a codec that rejects it rejects the array, but a codec that + accepts it has not thereby accepted every other chunk. """ if isinstance(chunk_grid, RegularChunkGridMetadata): return chunk_grid.chunk_shape return tuple(s if isinstance(s, int) else max(s) for s in chunk_grid.chunk_shapes) -def _distinct_chunk_shapes(chunk_grid: ChunkGridMetadata) -> Iterator[tuple[int, ...]]: - """Every distinct chunk shape occurring in `chunk_grid`. +def _defining_class(cls: type, name: str) -> type: + return next(klass for klass in cls.__mro__ if name in klass.__dict__) - For a rectilinear grid every combination of per-dimension distinct edges - occurs as an actual chunk shape (each edge along one dimension meets each - edge along every other), so this is the full cross product. - """ - if isinstance(chunk_grid, RegularChunkGridMetadata): - yield chunk_grid.chunk_shape - return - per_dim = ( - (s,) if isinstance(s, int) else tuple(dict.fromkeys(s)) for s in chunk_grid.chunk_shapes - ) - yield from itertools.product(*per_dim) - -def _resolved_chunk_specs( - chunk_grid: ChunkGridMetadata, spec: ArraySpec, codecs: tuple[Codec, ...] -) -> Iterator[ArraySpec]: - """Stream chunk specs through a codec prefix without storing the cross product.""" - for shape in _distinct_chunk_shapes(chunk_grid): - current = replace(spec, shape=shape) - for codec in codecs: - current = codec.resolve_metadata(current) - yield current +def _declared_chunk_grid( + codec: Codec, + *, + shape: tuple[int, ...], + chunk_grid: ChunkGridMetadata, + resolved_spec: ArraySpec, +) -> tuple[tuple[int, ...], ChunkGridMetadata] | None: + """The whole-array geometry after `codec`, if the codec declares one we trust. + + A declaration (`BaseCodec.resolve_chunk_grid`) is distrusted, and `None` + returned, when: + + - `resolve_metadata` is overridden in a subclass of the class that + declared the grid, since the declaration was written for a different + resolver (e.g. a subclass of a dtype-only codec that reshapes); or + - the declared grid's representative chunk shape disagrees with the + shape `resolve_metadata` produced for the representative chunk, which + proves the declaration wrong for at least that chunk. + """ + cls = type(codec) + grid_owner = _defining_class(cls, "resolve_chunk_grid") + metadata_owner = _defining_class(cls, "resolve_metadata") + if metadata_owner is not grid_owner and issubclass(metadata_owner, grid_owner): + return None + declared = codec.resolve_chunk_grid(shape=shape, chunk_grid=chunk_grid) + if declared is None or representative_chunk_shape(declared[1]) != resolved_spec.shape: + return None + return declared def evolve_and_validate_codecs( @@ -174,75 +180,82 @@ def evolve_and_validate_codecs( chunk_grid: ChunkGridMetadata, chunk_spec: ArraySpec, ) -> tuple[Codec, ...]: - """Evolve and validate codecs while preserving factored chunk geometry. - - Codecs with the default identity metadata resolver retain the grid. - CastValue and ScaleOffset also retain its geometry; TransposeCodec permutes - the dimensions of both the shape and grid. These paths resolve one spec per - codec and validate the grid directly, without enumerating chunk combinations. - Only the exact built-in classes use specialized geometry rules; overridden - metadata resolvers must use the general path. - - Arbitrary metadata resolvers stream each distinct chunk spec. Validation - uses the array geometry until a chunk shape changes, then each resolved - shape is validated as a regular grid. Full specs retain per-chunk dtype and - fill values. Prefixes are replayed for subsequent codecs rather than cached: - this bounds memory but may still require combinatorial work for arbitrary - transforms. Evolution uses one representative spec to produce one chain. + """Evolve each codec against the chunk spec it sees, and validate it. + + Every codec is validated against the geometry produced by the codecs + before it: a sharding codec placed after a transpose must divide the + transposed chunks, not the array's chunk grid. That geometry is threaded + as a whole chunk grid, never by enumerating chunks, so the cost is + proportional to the number of codecs and dimensions regardless of how many + distinct chunk shapes a rectilinear grid has. The approach mirrors zarrs, + whose array-to-array codecs map a chunk grid to a chunk grid and fall back + to "chunk-local" geometry when no whole grid exists. + + After each codec, the grid is carried forward in one of three ways: + + 1. **Declared.** The codec's `resolve_chunk_grid` maps the grid (identity + for codecs that do not change chunk shape, a permutation for + `TransposeCodec`). Validation stays exact: downstream codecs see every + chunk size through the grid. + 2. **Regular.** The codec does not declare a grid, but the grid is regular, + so every chunk has the same shape. Resolving the one representative spec + is exact, and the result is again a regular grid. + 3. **Chunk-local.** The codec does not declare a grid and the grid is + rectilinear. The chunks may now have shapes that no single grid + describes, and determining them would mean resolving every combination + of per-dimension edges, a cost that grows as the product of the edge + counts (a 3-d grid with 1000 distinct edges per axis has 10^9). Instead, + this and every later codec is validated against the representative + chunk only, and the geometry stays chunk-local to the end of the chain. + + Trade-off of chunk-local validation: a rejection is always correct, since + the representative is a real chunk of the array, but an acceptance is + incomplete. A chain that is invalid for some *other* chunk shape (for + example, a sharding codec whose inner chunk size divides the largest chunk + but not a smaller one) is accepted when the metadata is created and fails + when a chunk with that shape is first encoded or decoded. Codecs whose + correctness depends on the chunk shape must therefore check it again at + run time, as `ShardingCodec` does; a codec that skips that check can + silently corrupt data. A codec author avoids the chunk-local path entirely + by implementing `resolve_chunk_grid`. + + Only rectilinear grids can become chunk-local; a chain on a regular grid is + always validated exactly. Validation precedes metadata resolution, whose implementation may require - the validated dtype and geometry. Each codec is evolved only once here. + the validated dtype and geometry. Evolution uses the representative spec, + producing one codec chain for all chunks. """ - from zarr.codecs.cast_value import CastValue - from zarr.codecs.scale_offset import ScaleOffset - from zarr.codecs.transpose import TransposeCodec - out: list[Codec] = [] spec = chunk_spec - grid = chunk_grid - validation_shape = shape - fallback_grid: ChunkGridMetadata | None = None - fallback_spec = chunk_spec - fallback_codecs: tuple[Codec, ...] = () - shapes_changed = False + grid: ChunkGridMetadata | None = chunk_grid for codec in codecs: evolved = codec.evolve_from_array_spec(spec) - preserves_shape = type(evolved).resolve_metadata is BaseCodec.resolve_metadata or type( - evolved - ) in (CastValue, ScaleOffset) - if fallback_grid is None and (preserves_shape or type(evolved) is TransposeCodec): - evolved.validate(shape=validation_shape, dtype=spec.dtype, chunk_grid=grid) + if grid is None: + # chunk-local: see "Chunk-local" above + evolved.validate( + shape=spec.shape, + dtype=spec.dtype, + chunk_grid=RegularChunkGridMetadata(chunk_shape=spec.shape), + ) spec = evolved.resolve_metadata(spec) - if type(evolved) is TransposeCodec: - order = evolved.order - validation_shape = tuple(validation_shape[d] for d in order) - if isinstance(grid, RegularChunkGridMetadata): - grid = RegularChunkGridMetadata( - chunk_shape=tuple(grid.chunk_shape[d] for d in order) - ) - else: - grid = RectilinearChunkGridMetadata( - chunk_shapes=tuple(grid.chunk_shapes[d] for d in order) - ) else: - if fallback_grid is None: - fallback_grid, fallback_spec = grid, spec - changed = False - for current in _resolved_chunk_specs(fallback_grid, fallback_spec, fallback_codecs): - evolved.validate( - shape=current.shape if shapes_changed else validation_shape, - dtype=current.dtype, - chunk_grid=( - RegularChunkGridMetadata(chunk_shape=current.shape) - if shapes_changed - else grid - ), - ) - resolved = evolved.resolve_metadata(current) - changed |= resolved.shape != current.shape - shapes_changed |= changed - fallback_codecs += (evolved,) - spec = evolved.resolve_metadata(spec) + evolved.validate(shape=shape, dtype=spec.dtype, chunk_grid=grid) + resolved = evolved.resolve_metadata(spec) + declared = _declared_chunk_grid( + evolved, shape=shape, chunk_grid=grid, resolved_spec=resolved + ) + if declared is not None: + shape, grid = declared + elif isinstance(grid, RegularChunkGridMetadata): + if resolved.shape != spec.shape: + shape, grid = ( + resolved.shape, + RegularChunkGridMetadata(chunk_shape=resolved.shape), + ) + else: + grid = None + spec = resolved out.append(evolved) return tuple(out) diff --git a/tests/test_codecs/test_codec_chain_validation.py b/tests/test_codecs/test_codec_chain_validation.py index aa628ad607..b90f843570 100644 --- a/tests/test_codecs/test_codec_chain_validation.py +++ b/tests/test_codecs/test_codec_chain_validation.py @@ -6,6 +6,13 @@ combining ``reshape`` with ``transpose`` to both reorder and reshape; the ``transpose`` order then refers to the *reshaped* rank, so validating it against the array-level shape must not reject the chain. + +Geometry is threaded through the chain as a whole chunk grid (see +`evolve_and_validate_codecs`). On a rectilinear grid, a codec that does not +declare its grid (`resolve_chunk_grid`) makes the rest of the chain +"chunk-local": only a representative chunk is validated when the metadata is +created, and other chunk shapes are checked when they are encoded or decoded. +The tests below pin both sides of that trade-off. """ from __future__ import annotations @@ -17,20 +24,25 @@ import pytest import zarr -from zarr.abc.codec import ArrayArrayCodec +from zarr.abc.codec import ArrayArrayCodec, Codec from zarr.codecs import BytesCodec, ShardingCodec, TransposeCodec +from zarr.codecs.numcodecs import Delta +from zarr.codecs.scale_offset import ScaleOffset +from zarr.core.array_spec import ArrayConfig, ArraySpec +from zarr.core.buffer import default_buffer_prototype from zarr.core.dtype import Int32 from zarr.core.metadata.v3 import ( ArrayV3Metadata, + ChunkGridMetadata, RectilinearChunkGridMetadata, RegularChunkGridMetadata, + _declared_chunk_grid, ) from zarr.registry import _codec_registries, register_codec if TYPE_CHECKING: from collections.abc import Iterator - from zarr.core.array_spec import ArraySpec from zarr.core.buffer import NDBuffer from zarr.core.common import JSON @@ -69,16 +81,47 @@ def compute_encoded_size(self, input_byte_length: int, _chunk_spec: ArraySpec) - return input_byte_length +@dataclass(frozen=True) +class OpaqueCodec(ArrayArrayCodec): + """A pass-through filter that overrides `resolve_metadata` but does not + declare its chunk grid, as a typical third-party codec would. Metadata + validation cannot tell that it preserves chunk shape.""" + + is_fixed_size = True + + @classmethod + def from_dict(cls, data: dict[str, JSON]) -> Self: + return cls() + + def to_dict(self) -> dict[str, JSON]: + return {"name": "opaque"} + + def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec: + return chunk_spec + + async def _decode_single(self, chunk_array: NDBuffer, chunk_spec: ArraySpec) -> NDBuffer: + return chunk_array + + async def _encode_single(self, chunk_array: NDBuffer, chunk_spec: ArraySpec) -> NDBuffer: + return chunk_array + + def compute_encoded_size(self, input_byte_length: int, _chunk_spec: ArraySpec) -> int: + return input_byte_length + + @pytest.fixture(autouse=True) -def _register_reshape() -> Iterator[None]: - previous = _codec_registries.get("reshape") - register_codec("reshape", ReshapeCodec) +def _register_test_codecs() -> Iterator[None]: + names = {"reshape": ReshapeCodec, "opaque": OpaqueCodec} + previous = {name: _codec_registries.get(name) for name in names} + for name, codec_cls in names.items(): + register_codec(name, codec_cls) try: yield finally: - _codec_registries.pop("reshape", None) - if previous is not None: - _codec_registries["reshape"] = previous + for name, registry in previous.items(): + _codec_registries.pop(name, None) + if registry is not None: + _codec_registries[name] = registry SHAPE = (4, 6, 8) @@ -151,110 +194,213 @@ def test_sharding_inner_chain_is_validated() -> None: _metadata((bad,), chunk_shape=SHAPE) -def _rectilinear_transpose_sharding_metadata(inner: tuple[int, int]) -> ArrayV3Metadata: - """Rectilinear grid (chunks (4,5) and (6,5)), transposed, then sharded.""" +def _rectilinear_sharding_metadata( + filters: tuple[Any, ...], inner: tuple[int, int] +) -> ArrayV3Metadata: + """Rectilinear grid (chunks (4,5) and (6,5)), `filters`, then sharded.""" return ArrayV3Metadata( shape=(10, 5), data_type=Int32(), chunk_grid=RectilinearChunkGridMetadata(chunk_shapes=((4, 6), 5)), chunk_key_encoding={"name": "default"}, fill_value=0, - codecs=(TransposeCodec(order=(1, 0)), ShardingCodec(chunk_shape=inner)), + codecs=(*filters, ShardingCodec(chunk_shape=inner)), attributes=None, dimension_names=None, ) +TRANSPOSE = (TransposeCodec(order=(1, 0)),) +DELTA_TRANSPOSE = (Delta(dtype=" None: """After a shape-changing codec on a rectilinear grid, an inner shard shape dividing every transposed chunk shape ((5,4) and (5,6)) is accepted.""" with zarr.config.set({"array.rectilinear_chunks": True}): - meta = _rectilinear_transpose_sharding_metadata(inner) + meta = _rectilinear_sharding_metadata(filters, inner) assert ArrayV3Metadata.from_dict(meta.to_dict()) == meta -def test_rectilinear_every_chunk_shape_validated() -> None: - """Under a rectilinear grid, size-sensitive validation after a - shape-changing codec must consider every distinct chunk shape, not a single - representative: an inner shard size dividing the largest transposed chunk - (5,6) but not the smaller (5,4) is rejected.""" +@pytest.mark.parametrize("filters", [TRANSPOSE, DELTA_TRANSPOSE], ids=["transpose", "delta"]) +def test_rectilinear_every_chunk_shape_validated(filters: tuple[Any, ...]) -> None: + """When every codec before the sharding codec declares its chunk grid, + validation covers every distinct chunk shape, not a single representative: + an inner shard size dividing the largest transposed chunk (5,6) but not the + smaller (5,4) is rejected when the metadata is created.""" with ( zarr.config.set({"array.rectilinear_chunks": True}), pytest.raises(ValueError, match="not\\s+divisible"), ): - _rectilinear_transpose_sharding_metadata((5, 3)) + _rectilinear_sharding_metadata(filters, (5, 3)) -@pytest.mark.parametrize("rank", [4, 12]) -@pytest.mark.parametrize("chain", ["bytes", "scale-offset", "transpose-sharding"]) -def test_factored_validation_does_not_visit_chunk_combinations( - rank: int, chain: str, monkeypatch: pytest.MonkeyPatch -) -> None: - """Common codec chains resolve specs independently of the chunk cross product.""" - from zarr.abc.codec import BaseCodec - from zarr.codecs.scale_offset import ScaleOffset +def test_chunk_local_rejects_invalid_representative_chunk() -> None: + """After an undeclared codec on a rectilinear grid, the representative + chunk (the largest edge per axis, (6,5)) is still validated when the + metadata is created, so a shard size that fails it is rejected.""" + with ( + zarr.config.set({"array.rectilinear_chunks": True}), + pytest.raises(ValueError, match="not\\s+divisible"), + ): + _rectilinear_sharding_metadata((OpaqueCodec(),), (4, 5)) - calls = 0 - original = BaseCodec.resolve_metadata - - def counted(self: Any, spec: ArraySpec) -> ArraySpec: - nonlocal calls - calls += 1 - assert calls <= 10, "metadata validation enumerated chunk combinations" - return original(self, spec) - - monkeypatch.setattr(BaseCodec, "resolve_metadata", counted) - codecs: tuple[Any, ...] - if chain == "bytes": - codecs = (BytesCodec(),) - elif chain == "scale-offset": - codecs = (ScaleOffset(offset=0), BytesCodec()) - else: - codecs = ( - TransposeCodec(order=tuple(reversed(range(rank)))), - ShardingCodec(chunk_shape=(1,) * rank), + +@pytest.mark.filterwarnings( + "ignore:Combining a `sharding_indexed` codec:zarr.errors.ZarrUserWarning" +) +def test_chunk_local_defers_other_shard_shapes_to_run_time() -> None: + """The documented cost of chunk-local validation: after an undeclared codec + on a rectilinear grid, an inner shard size dividing the representative + chunk (6,5) but not the smaller (4,5) is accepted when the array is created. + The non-dividing shard is rejected when it is first encoded, instead of + being floor-divided and silently written with the wrong layout, and the + dividing shard still round-trips.""" + with zarr.config.set({"array.rectilinear_chunks": True}): + array = zarr.create_array( + {}, + shape=(10, 5), + chunks=[[4, 6], [5]], + dtype="i4", + filters=[OpaqueCodec()], + serializer=ShardingCodec(chunk_shape=(3, 5)), + compressors=None, ) + array[4:] = 1 + np.testing.assert_array_equal(array[4:], np.ones((6, 5), dtype="i4")) + with pytest.raises(ValueError, match="not divisible by the shard's inner chunk shape"): + array[:4] = 1 + + +def test_chunk_local_does_not_enumerate_to_find_codec_errors() -> None: + """A shape-changing undeclared codec that accepts the representative chunk + ((2,2) -> (4,)) but would reject a smaller one ((1,1)) is accepted when the + metadata is created: validation does not enumerate the chunks of the grid + to find the failing one. The codec is left to reject that chunk itself when + it is encoded or decoded.""" with zarr.config.set({"array.rectilinear_chunks": True}): ArrayV3Metadata( - shape=(3,) * rank, + shape=(3, 3), data_type=Int32(), - chunk_grid=RectilinearChunkGridMetadata(chunk_shapes=((1, 2),) * rank), + chunk_grid=RectilinearChunkGridMetadata(chunk_shapes=((1, 2), (1, 2))), chunk_key_encoding={"name": "default"}, fill_value=0, - codecs=codecs, + codecs=(ReshapeCodec(shape=(4,)), BytesCodec()), attributes={}, dimension_names=None, ) - assert calls <= 10 -def test_arbitrary_shape_validation_stops_at_first_invalid_chunk( - monkeypatch: pytest.MonkeyPatch, +class _ReshapingDelta(Delta): + """Overrides the resolver of a codec that declares an identity grid.""" + + def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec: + return replace(chunk_spec, shape=chunk_spec.shape[::-1]) + + +class _MisdeclaredCodec(OpaqueCodec): + """Declares an identity grid while `resolve_metadata` reverses the axes.""" + + def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec: + return replace(chunk_spec, shape=chunk_spec.shape[::-1]) + + def resolve_chunk_grid( + self, *, shape: tuple[int, ...], chunk_grid: ChunkGridMetadata + ) -> tuple[tuple[int, ...], ChunkGridMetadata]: + return shape, chunk_grid + + +@pytest.mark.parametrize( + ("codec", "expected"), + [ + (BytesCodec(), ((10, 5), ((4, 6), 5))), + (Delta(dtype=" None: - """A failing arbitrary transform does not consume all chunk combinations.""" - import itertools + """A codec's declared grid is used only when it is trustworthy: declared + by the class whose resolver is in effect, and consistent with that + resolver on the representative chunk.""" + spec = ArraySpec( + shape=(6, 5), + dtype=Int32(), + fill_value=0, + config=ArrayConfig.from_dict({}), + prototype=default_buffer_prototype(), + ) + with zarr.config.set({"array.rectilinear_chunks": True}): + grid = RectilinearChunkGridMetadata(chunk_shapes=((4, 6), 5)) + declared = _declared_chunk_grid( + codec, shape=(10, 5), chunk_grid=grid, resolved_spec=codec.resolve_metadata(spec) + ) + if expected is None: + assert declared is None + else: + shape, chunk_shapes = expected + assert declared == (shape, RectilinearChunkGridMetadata(chunk_shapes=chunk_shapes)) - original_product = itertools.product - def guarded_product(*args: Any) -> Iterator[tuple[int, ...]]: - yield next(original_product(*args)) - raise AssertionError("consumed combinations after the first invalid chunk") +CHAINS: dict[str, Any] = { + "bytes": lambda rank: (BytesCodec(),), + "scale-offset": lambda rank: (ScaleOffset(offset=0), BytesCodec()), + "delta": lambda rank: (Delta(dtype=" None: + """Validation cost is independent of the number of chunk shapes. The grid + has 2**rank distinct chunk shapes (4096 at rank 12); resolving metadata once + per codec, rather than once per chunk shape, keeps the call count small for + declared, default, and undeclared codecs alike.""" + codecs = CHAINS[chain](rank) + calls = 0 + + def count(resolver: Any) -> Any: + def counted(self: Any, spec: ArraySpec) -> ArraySpec: + nonlocal calls + calls += 1 + assert calls <= 10, "metadata validation enumerated chunk combinations" + return cast("ArraySpec", resolver(self, spec)) + + return counted + + # Patch each resolver on the class that defines it, so that which class + # defines `resolve_metadata` (and hence which grid declaration is trusted) + # is unchanged. + owners = {next(k for k in type(c).__mro__ if "resolve_metadata" in k.__dict__) for c in codecs} + for owner in owners: + monkeypatch.setattr(owner, "resolve_metadata", count(owner.__dict__["resolve_metadata"])) + with zarr.config.set({"array.rectilinear_chunks": True}): ArrayV3Metadata( - shape=(3, 3), + shape=(3,) * rank, data_type=Int32(), - chunk_grid=RectilinearChunkGridMetadata(chunk_shapes=((1, 2), (1, 2))), + chunk_grid=RectilinearChunkGridMetadata(chunk_shapes=((1, 2),) * rank), chunk_key_encoding={"name": "default"}, fill_value=0, - codecs=(ReshapeCodec(shape=(4,)), BytesCodec()), + codecs=codecs, attributes={}, dimension_names=None, ) + assert 0 < calls <= 10 From 55c8bfadd89ad89b2ce7e55969f52d02c0319477 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 14 Sep 2026 17:33:54 +0200 Subject: [PATCH 11/11] fix(codecs): evolve rectilinear pipelines against the representative chunk shape `create_codec_pipeline` evolved V3 pipelines against an all-ones placeholder chunk spec whenever the grid was not regular, so any codec whose `resolve_metadata` depends on the chunk shape (e.g. a reshape filter) failed at array creation for rectilinear grids even though metadata validation had accepted the chain, and shape-sensitive evolution output (sharding's inner chain, BytesCodec endian) was computed against a meaningless shape. Use `representative_chunk_shape`, the same shape `ArrayV3Metadata.__init__` threads through evolution, so the pipeline carries exactly the codecs the metadata produced. The ChunkTransform itself remains shape-agnostic: it only holds the evolved codecs and resolves specs per call. Assisted-by: ClaudeCode:claude-fable-5-1 --- changes/4352.bugfix.md | 2 ++ src/zarr/core/array.py | 17 +++++----- .../test_codec_chain_validation.py | 32 +++++++++++++++++++ 3 files changed, 43 insertions(+), 8 deletions(-) diff --git a/changes/4352.bugfix.md b/changes/4352.bugfix.md index c52e343521..73fd5acafe 100644 --- a/changes/4352.bugfix.md +++ b/changes/4352.bugfix.md @@ -3,3 +3,5 @@ Validate each codec in a chain against the chunk geometry produced by the codecs Geometry is threaded through the chain as a whole chunk grid, so validation cost no longer depends on the number of distinct chunk shapes. Codecs describe how they map the grid with the new `BaseCodec.resolve_chunk_grid` method; the built-in dtype and fill-value codecs (`cast_value`, `scale_offset`, and the numcodecs `delta`, `fixedscaleoffset` and `astype` filters) declare the identity and `transpose` declares a permutation. Trade-off: on a rectilinear chunk grid, a codec that overrides `resolve_metadata` without implementing `resolve_chunk_grid` makes the rest of its chain "chunk-local". Codecs after it are validated against one representative chunk (the largest edge along each axis), so a chain that is invalid only for some other chunk shape is accepted when the array is created and fails when such a chunk is first written or read. `ShardingCodec` now checks divisibility at encode and decode time, so this case raises an error instead of corrupting data. Regular chunk grids are always validated exactly. See "Chunk geometry and validation" in the extending guide. + +The codec pipeline of a rectilinear-chunked array is now evolved against that same representative chunk shape, instead of an all-ones placeholder that broke shape-dependent codecs at array creation even when the metadata validated. diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 047e7bb3b3..6ea1fde5ad 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -131,6 +131,7 @@ RegularChunkGridMetadata, create_chunk_grid_metadata, parse_node_type_array, + representative_chunk_shape, ) from zarr.core.sync import sync from zarr.errors import ( @@ -239,15 +240,15 @@ def create_codec_pipeline(metadata: ArrayMetadata, *, store: Store | None = None # `codecs_from_list_unchecked`, so it does not re-emit them. pipeline = get_pipeline_class().from_codecs(metadata.codecs) - # Use the regular chunk shape if available, otherwise use a - # placeholder. The ChunkTransform is shape-agnostic — the actual - # chunk shape is passed per-call at decode/encode time. - if isinstance(metadata.chunk_grid, RegularChunkGridMetadata): - chunk_shape = metadata.chunk_grid.chunk_shape - else: - chunk_shape = (1,) * len(metadata.shape) + # Evolve against the same representative chunk shape that + # `ArrayV3Metadata.__init__` used, so a codec whose `resolve_metadata` + # depends on the chunk shape (reshape, sharding's inner chain) sees a + # real chunk rather than a placeholder that metadata validation never + # saw. Only evolution is shape-sensitive: the resulting ChunkTransform + # is shape-agnostic — the actual chunk shape is passed per call at + # decode/encode time. chunk_spec = ArraySpec( - shape=chunk_shape, + shape=representative_chunk_shape(metadata.chunk_grid), dtype=metadata.data_type, fill_value=metadata.fill_value, config=ArrayConfig.from_dict({}), diff --git a/tests/test_codecs/test_codec_chain_validation.py b/tests/test_codecs/test_codec_chain_validation.py index b90f843570..04dff7745b 100644 --- a/tests/test_codecs/test_codec_chain_validation.py +++ b/tests/test_codecs/test_codec_chain_validation.py @@ -17,6 +17,7 @@ from __future__ import annotations +from collections.abc import Iterable from dataclasses import dataclass, replace from typing import TYPE_CHECKING, Any, Self, cast @@ -153,6 +154,37 @@ def test_rank_changing_chain_roundtrip(shards: tuple[int, ...] | None) -> None: assert np.array_equal(reloaded[:], data) +@pytest.mark.parametrize( + "chunks", + [(2, 4), [[2, 2], 4], [[2, 2], [4]]], + ids=["regular", "rectilinear-mixed", "rectilinear-explicit"], +) +def test_pipeline_evolves_against_representative_chunk( + chunks: tuple[int, ...] | list[Any], +) -> None: + """The codec pipeline is evolved against the same representative chunk shape + that metadata validation used, for regular and rectilinear grids alike. A + reshape whose size matches every (2, 4) chunk but not a placeholder chunk + must be accepted at array creation, and the evolved pipeline must carry the + codecs metadata validation produced.""" + data = np.arange(16, dtype="i4").reshape(4, 4) + with zarr.config.set({"array.rectilinear_chunks": True}): + a = zarr.create_array( + {}, + shape=(4, 4), + chunks=chunks, + dtype="i4", + filters=[ReshapeCodec(shape=(8,))], + compressors=None, + ) + pipeline = a._async_array.codec_pipeline + assert isinstance(pipeline, Iterable) + assert isinstance(a.metadata, ArrayV3Metadata) + assert tuple(pipeline) == a.metadata.codecs + a[:] = data + assert np.array_equal(a[:], data) + + def _metadata(codecs: tuple[Any, ...], chunk_shape: tuple[int, ...] = CHUNKS) -> ArrayV3Metadata: return ArrayV3Metadata( shape=SHAPE,