From 27f60ad6b14ddf1d15283e76366ee60df91c33fe Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 9 Sep 2026 12:57:46 +0200 Subject: [PATCH 1/3] test: give the sharded coordinate-selection write path real weight in test_oindex Add zarr.testing.strategies.sharded_arrays, a strategy that always yields a sharded v3 array (drawing chunk shape, shard shape, subchunk write order and inner codec chain), and draw it as a third arm in test_oindex and test_vindex. simple_arrays shards only a few percent of its draws, so under the derandomized ci profile test_oindex reached the ShardingCodec value-reshape guard (GH4284/GH4316) in 2 of 300 examples; it now reaches it in about 50 (13-27% of writes across random seeds), and the test fails without the guard. The sharded arm in test_oindex draws at least two dimensions, since the guard only applies to a value with two or more axes. A Hypothesis event at the write site reports the reach under --hypothesis-show-statistics. Narrow the docstring of test_set_selection_rejects_value_with_wrong_rank to the cases it pins: storage layout can change which writes are rejected, e.g. oindex[[3, 1], [0, 2]] with a (2, 2, 1) value is accepted chunked but raises sharded. Assisted-by: ClaudeCode:claude-fable-5-1 --- changes/0000.misc.md | 1 + src/zarr/testing/strategies.py | 70 ++++++++++++++++++++++++++++++---- tests/test_indexing.py | 15 +++++--- tests/test_properties.py | 19 ++++++++- 4 files changed, 91 insertions(+), 14 deletions(-) create mode 100644 changes/0000.misc.md diff --git a/changes/0000.misc.md b/changes/0000.misc.md new file mode 100644 index 0000000000..4ae316b0a1 --- /dev/null +++ b/changes/0000.misc.md @@ -0,0 +1 @@ +Added `zarr.testing.strategies.sharded_arrays`, a Hypothesis strategy that always generates a sharded Zarr v3 array, drawing the inner chunk shape, shard shape, subchunk write order and inner codec chain. The `test_oindex` and `test_vindex` property tests now draw it as a third arm alongside `simple_arrays` and `rectilinear_arrays`, so the sharding codec's write path for orthogonal selections with two or more array-indexed axes is exercised in tens of examples per run instead of about one. diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index dfdf25eafd..c8460f3584 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -255,6 +255,19 @@ def shard_shapes( return tuple(m * c for m, c in zip(multiples, chunk_shape, strict=True)) +@st.composite +def _sharding_codecs(draw: st.DrawFn, *, chunk_shape: tuple[int, ...]) -> ShardingCodec: + """A ``ShardingCodec`` over ``chunk_shape`` with a drawn subchunk write order and inner codec chain.""" + subchunk_write_order = draw(subchunk_write_orders) + inner_codecs = draw(sharding_inner_codecs, label="sharding inner codecs") + return ShardingCodec( + subchunk_write_order=subchunk_write_order, + codecs=inner_codecs, + index_codecs=[BytesCodec(), Crc32cCodec()], + chunk_shape=chunk_shape, + ) + + @st.composite def np_array_and_chunks( draw: st.DrawFn, @@ -334,14 +347,7 @@ def arrays( ) event("sharded" if shard_shape is not None else "unsharded") if shard_shape is not None: - subchunk_write_order = draw(subchunk_write_orders) - inner_codecs = draw(sharding_inner_codecs, label="sharding inner codecs") - serializer = ShardingCodec( - subchunk_write_order=subchunk_write_order, - codecs=inner_codecs, - index_codecs=[BytesCodec(), Crc32cCodec()], - chunk_shape=chunks_param, - ) + serializer = draw(_sharding_codecs(chunk_shape=chunks_param)) compressors_unsearched = None else: chunks_param = draw(chunk_shapes(shape=nparray.shape), label="chunk shape") @@ -533,6 +539,54 @@ def rectilinear_arrays( return a +# Sharded arrays need min_side >= 1: a shard must hold at least one chunk on every axis. +_sharded_shapes = npst.array_shapes(max_dims=4, min_side=1, max_side=8) + + +@st.composite +def sharded_arrays( + draw: st.DrawFn, + *, + shapes: st.SearchStrategy[tuple[int, ...]] = _sharded_shapes, +) -> Any: + """Generate a zarr v3 array whose chunks are grouped into shards. + + ``arrays`` shards only a small fraction of its draws (a v3 array with a + regular chunk grid, every axis larger than a chunk that is itself larger + than 1, and then only half the time), so a property test that must + exercise the sharding codec should draw from this strategy directly. Every + draw is sharded: the inner chunk shape and the shard shape (an integral + number of chunks per axis, possibly a single chunk) are drawn from + ``shapes``, and the codec's subchunk write order and inner codec chain are + drawn as in ``arrays``. ``shapes`` must generate shapes with at least one + element on every axis. + """ + shape = draw(shapes) + chunk_shape = draw(chunk_shapes(shape=shape), label="chunk shape") + shard_shape = draw(shard_shapes(shape=shape, chunk_shape=chunk_shape), label="shard shape") + serializer = draw(_sharding_codecs(chunk_shape=chunk_shape)) + + nparray = draw(numpy_arrays(shapes=st.just(shape)), label="array data") + fill_value = draw(st.one_of([st.none(), npst.from_dtype(nparray.dtype)])) + dim_names = draw(dimension_names(ndim=len(shape)), label="dimension names") + + a = zarr.create_array( + store=MemoryStore(), + shape=shape, + chunks=chunk_shape, + shards=shard_shape, + dtype=nparray.dtype, + fill_value=fill_value, + dimension_names=dim_names, + serializer=serializer, + compressors=None, + ) + assert a.shards == shard_shape + assert a.chunks == chunk_shape + a[:] = nparray + return a + + def is_negative_slice(idx: Any) -> bool: return isinstance(idx, slice) and idx.step is not None and idx.step < 0 diff --git a/tests/test_indexing.py b/tests/test_indexing.py index 4a1a4172d2..54de72e507 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -2135,14 +2135,19 @@ def test_set_selection_rejects_value_with_wrong_rank( shards: tuple[int, ...] | None, pipeline_path: str, ) -> None: - """A value whose rank does not fit the selection raises regardless of storage layout. + """These wrong-rank values are rejected on chunked and sharded arrays alike. The sharding codec re-derives an indexer from the selection it is handed and ravels the value when it is the selection's broadcast shape minus - integer-indexed axes. Any other rank must fail on a sharded array exactly - as it does on a chunked one; an element count that happens to match is - not grounds to accept it. Only the rejection is asserted: a write that - fails inside the chunk merge may already have touched other chunks. + integer-indexed axes. The cases here pin that a matching element count + alone does not make the codec accept a value the chunked path rejects. + That is not a general law: storage layout can change which writes are + rejected. ``oindex[np.array([3, 1]), np.array([0, 2])]`` with a + ``(2, 2, 1)`` value is accepted on a chunked ``(4, 4)`` array with + ``(2, 2)`` chunks, because each chunk receives a ``(1, 1, 1)`` piece numpy + can broadcast, while the sharded array raises ``ValueError`` and numpy + rejects it outright. Only the rejection is asserted: a write that fails + inside the chunk merge may already have touched other chunks. """ a = np.zeros((4, 4), dtype=np.int32) value = np.arange(np.prod(value_shape), dtype=np.int32).reshape(value_shape) diff --git a/tests/test_properties.py b/tests/test_properties.py index 2b82d54e2e..ee30e6efe0 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -15,7 +15,7 @@ import hypothesis.extra.numpy as npst import hypothesis.strategies as st -from hypothesis import assume, given, settings +from hypothesis import assume, event, given, settings from zarr.abc.store import Store from zarr.core.common import ZARR_JSON, ZARRAY_JSON, ZATTRS_JSON @@ -31,6 +31,7 @@ numpy_arrays, orthogonal_indices, rectilinear_arrays, + sharded_arrays, simple_arrays, stores, zarr_formats, @@ -158,10 +159,17 @@ async def test_basic_indexing_complex_rectilinear(data: st.DataObject) -> None: @pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") async def test_oindex(data: st.DataObject) -> None: # integer_array_indices can't handle 0-size dimensions. + # A sharded array is drawn as its own arm: simple_arrays shards only a few + # percent of its draws, and the sharding codec's write path for a selection + # with two or more array-indexed axes (GH4284) needs real weight here. That + # path only exists for a value with two or more axes, hence min_dims=2. zarray = data.draw( st.one_of( simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1)), rectilinear_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1, max_side=20)), + sharded_arrays( + shapes=npst.array_shapes(min_dims=2, max_dims=4, min_side=1, max_side=8) + ), ) ) nparray = zarray[:] @@ -181,6 +189,14 @@ async def test_oindex(data: st.DataObject) -> None: if isinstance(idxr, np.ndarray) and idxr.size != np.unique(idxr).size: # behaviour of setitem with repeated indices is not guaranteed in practice assume(False) + # The sharding codec sees a coordinate selection (the GH4284 path) when the + # chunk selection has more than one array axis or drops an integer axis. + n_array_axes = sum(isinstance(idxr, np.ndarray) for idxr in zindexer) + coordinate_path = n_array_axes > 1 or any(isinstance(idxr, int) for idxr in zindexer) + event( + f"oindex write: {'sharded' if zarray.shards is not None else 'unsharded'}, " + f"{'coordinate' if coordinate_path else 'orthogonal'} chunk selection" + ) new_data = data.draw(numpy_arrays(shapes=st.just(actual.shape), dtype=nparray.dtype)) nparray[npindexer] = new_data zarray.oindex[zindexer] = new_data @@ -197,6 +213,7 @@ async def test_vindex(data: st.DataObject) -> None: st.one_of( simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1)), rectilinear_arrays(shapes=npst.array_shapes(max_dims=3, min_side=1, max_side=20)), + sharded_arrays(), ) ) nparray = zarray[:] From 06e3f00ca46cb2dad866dfdded50d1056378fd52 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 9 Sep 2026 13:03:45 +0200 Subject: [PATCH 2/3] chore: rename changelog fragment to the PR number Assisted-by: ClaudeCode:claude-fable-5-1 --- changes/{0000.misc.md => 4331.misc.md} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename changes/{0000.misc.md => 4331.misc.md} (100%) diff --git a/changes/0000.misc.md b/changes/4331.misc.md similarity index 100% rename from changes/0000.misc.md rename to changes/4331.misc.md From 67e69cd7248504dcd8252206ce94e903a86c34ab Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 9 Sep 2026 15:15:57 +0200 Subject: [PATCH 3/3] test(strategies): let sharded_arrays nest one level of sharding Review asked for recursive sharding in the new strategy; one level is enough to reach the shard-within-shard paths, which is where the original GH4280 failure lived. Building it exposed that the strategy already nested by accident: passing `shards=` makes `create_array` wrap the drawn `ShardingCodec` in a second one with the same chunk shape, so every "sharded" draw was a shard of one-chunk inner shards and the drawn subchunk write order only ever reached the inner codec, never the outer. `sharded_arrays` now makes the shard the array's chunk grid and passes the drawn codec as the serializer, so a single-level draw really is single-level and the outer codec carries the drawn write order. A drawn (or `nested=`) flag then wraps that codec in another `ShardingCodec` whose chunk shape is an integral number of chunks and whose shard is an integral number of those. `_sharding_codecs` takes an optional inner codec chain to build the outer layer. (`arrays()` still uses the `shards=` form and so keeps the accidental one-chunk nesting; left alone here.) Measured on test_oindex under the ci profile: 24% of writes take the sharded coordinate-selection path, 16% of examples are nested, and the outer codec sees all four write orders instead of only `morton`. Assisted-by: ClaudeCode:claude-fable-5-1 --- changes/4331.misc.md | 2 +- src/zarr/testing/strategies.py | 67 +++++++++++++++++++++++++++------- 2 files changed, 55 insertions(+), 14 deletions(-) diff --git a/changes/4331.misc.md b/changes/4331.misc.md index 4ae316b0a1..949ebfffd3 100644 --- a/changes/4331.misc.md +++ b/changes/4331.misc.md @@ -1 +1 @@ -Added `zarr.testing.strategies.sharded_arrays`, a Hypothesis strategy that always generates a sharded Zarr v3 array, drawing the inner chunk shape, shard shape, subchunk write order and inner codec chain. The `test_oindex` and `test_vindex` property tests now draw it as a third arm alongside `simple_arrays` and `rectilinear_arrays`, so the sharding codec's write path for orthogonal selections with two or more array-indexed axes is exercised in tens of examples per run instead of about one. +Added `zarr.testing.strategies.sharded_arrays`, a Hypothesis strategy that always generates a sharded Zarr v3 array, drawing the chunk shape, shard shape, subchunk write order and inner codec chain, and — for half of its draws, or as selected by its `nested` argument — one level of recursive sharding, where the chunks are grouped into inner shards that are in turn grouped into the stored shards. The `test_oindex` and `test_vindex` property tests now draw it as a third arm alongside `simple_arrays` and `rectilinear_arrays`, so the sharding codec's write path for orthogonal selections with two or more array-indexed axes is exercised in tens of examples per run instead of about one. diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index c8460f3584..db01697f1e 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -1,7 +1,7 @@ import itertools import math import sys -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence from typing import Any, Literal import hypothesis.extra.numpy as npst @@ -12,6 +12,7 @@ from hypothesis.strategies import SearchStrategy import zarr +from zarr.abc.codec import Codec from zarr.abc.store import ( ByteRequest, OffsetByteRequest, @@ -256,10 +257,21 @@ def shard_shapes( @st.composite -def _sharding_codecs(draw: st.DrawFn, *, chunk_shape: tuple[int, ...]) -> ShardingCodec: - """A ``ShardingCodec`` over ``chunk_shape`` with a drawn subchunk write order and inner codec chain.""" +def _sharding_codecs( + draw: st.DrawFn, + *, + chunk_shape: tuple[int, ...], + codecs: Sequence[Codec] | None = None, +) -> ShardingCodec: + """A ``ShardingCodec`` over ``chunk_shape`` with a drawn subchunk write order. + + The inner codec chain is drawn from ``sharding_inner_codecs`` unless ``codecs`` + is given, which lets a caller nest another ``ShardingCodec`` inside. + """ subchunk_write_order = draw(subchunk_write_orders) - inner_codecs = draw(sharding_inner_codecs, label="sharding inner codecs") + inner_codecs: Sequence[Codec] = ( + draw(sharding_inner_codecs, label="sharding inner codecs") if codecs is None else codecs + ) return ShardingCodec( subchunk_write_order=subchunk_write_order, codecs=inner_codecs, @@ -548,6 +560,7 @@ def sharded_arrays( draw: st.DrawFn, *, shapes: st.SearchStrategy[tuple[int, ...]] = _sharded_shapes, + nested: bool | None = None, ) -> Any: """Generate a zarr v3 array whose chunks are grouped into shards. @@ -555,34 +568,62 @@ def sharded_arrays( regular chunk grid, every axis larger than a chunk that is itself larger than 1, and then only half the time), so a property test that must exercise the sharding codec should draw from this strategy directly. Every - draw is sharded: the inner chunk shape and the shard shape (an integral - number of chunks per axis, possibly a single chunk) are drawn from - ``shapes``, and the codec's subchunk write order and inner codec chain are - drawn as in ``arrays``. ``shapes`` must generate shapes with at least one - element on every axis. + draw is sharded: the chunk shape and the shard shape (an integral number of + chunks per axis, possibly a single chunk) are drawn from ``shapes``, and + the codec's subchunk write order and inner codec chain are drawn as in + ``arrays``. ``shapes`` must generate shapes with at least one element on + every axis. + + ``nested`` selects one level of recursive sharding: the drawn chunks are + grouped into inner shards, which are themselves grouped into the shards + stored in the array, so the outer ``ShardingCodec`` wraps an inner one with + its own subchunk write order. ``None`` (the default) draws it, so half the + examples nest. For a nested array ``Array.chunks`` is the inner shard shape + (the outer codec's chunk shape); the innermost chunk shape is the inner + codec's ``chunk_shape``. """ shape = draw(shapes) chunk_shape = draw(chunk_shapes(shape=shape), label="chunk shape") - shard_shape = draw(shard_shapes(shape=shape, chunk_shape=chunk_shape), label="shard shape") serializer = draw(_sharding_codecs(chunk_shape=chunk_shape)) + nest = draw(st.booleans(), label="nested sharding") if nested is None else nested + if nest: + # Each level's shard is an integral number of the level below's chunks. + codec_chunk_shape = draw( + shard_shapes(shape=shape, chunk_shape=chunk_shape), label="inner shard shape" + ) + serializer = draw(_sharding_codecs(chunk_shape=codec_chunk_shape, codecs=[serializer])) + else: + codec_chunk_shape = chunk_shape + shard_shape = draw( + shard_shapes(shape=shape, chunk_shape=codec_chunk_shape), label="shard shape" + ) + event("nested sharding" if nest else "single-level sharding") nparray = draw(numpy_arrays(shapes=st.just(shape)), label="array data") fill_value = draw(st.one_of([st.none(), npst.from_dtype(nparray.dtype)])) dim_names = draw(dimension_names(ndim=len(shape)), label="dimension names") + # The shard is the array's chunk grid and the drawn codec is its serializer. + # Passing ``shards=`` instead would make ``create_array`` wrap the codec in a + # second ``ShardingCodec`` of the same chunk shape, hiding the drawn write + # order behind a default outer one. a = zarr.create_array( store=MemoryStore(), shape=shape, - chunks=chunk_shape, - shards=shard_shape, + chunks=shard_shape, dtype=nparray.dtype, fill_value=fill_value, dimension_names=dim_names, serializer=serializer, + filters=None, compressors=None, ) assert a.shards == shard_shape - assert a.chunks == chunk_shape + assert a.chunks == codec_chunk_shape + assert isinstance(a.metadata, ArrayV3Metadata) + (codec,) = a.metadata.codecs + assert isinstance(codec, ShardingCodec) + assert codec.subchunk_write_order == serializer.subchunk_write_order a[:] = nparray return a