From 0a9599db9e2817b8ebe95ae0a4337a97217852a5 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 17 Aug 2026 17:47:09 +0200 Subject: [PATCH 1/3] 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 41780e45b4..a25d1bbb2d 100644 --- a/src/zarr/codecs/sharding.py +++ b/src/zarr/codecs/sharding.py @@ -69,6 +69,7 @@ ChunkGridMetadata, RectilinearChunkGridMetadata, RegularChunkGridMetadata, + evolve_and_validate_codecs, parse_codecs, ) from zarr.registry import get_ndbuffer_class, get_pipeline_class @@ -593,6 +594,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 fc47f8fc95..bf82182fd7 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -113,6 +113,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 @@ -503,28 +561,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) @@ -541,8 +594,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): @@ -551,8 +604,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 d54ffe886d207a7e5d5865378c037da8c798c840 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 30 Aug 2026 16:20:49 +0200 Subject: [PATCH 2/3] 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 | 94 +++++++++++++++---- .../test_codec_chain_validation.py | 32 ++++++- 2 files changed, 108 insertions(+), 18 deletions(-) diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index bf82182fd7..1ad4ae9077 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,12 @@ from zarr.core.dtype.common import check_dtype_spec_v3 from zarr.core.json_parse import parse_field, validate_json_value from zarr.core.metadata.common import parse_attributes -from zarr.errors import MetadataValidationError, NodeTypeValidationError, UnknownCodecError +from zarr.errors import ( + MetadataValidationError, + NodeTypeValidationError, + UnknownCodecError, + ZarrUserWarning, +) from zarr.registry import get_codec_class if TYPE_CHECKING: @@ -117,15 +124,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], *, @@ -145,11 +180,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 @@ -157,17 +200,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 a9f4a27876e79582369a3be3911a7c83a74a968f Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 30 Aug 2026 19:15:21 +0200 Subject: [PATCH 3/3] 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()