test: give the sharded coordinate-selection write path real weight in test_oindex (#4331)

* 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

* chore: rename changelog fragment to the PR number

Assisted-by: ClaudeCode:claude-fable-5-1

* 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

* chore: strip trailing whitespace from the #3285 changelog fragment

Introduced on main by #4330; it fails pre-commit.ci's trailing-whitespace
hook on every PR that merges main from now on. Fixed here so this PR is
green and main is fixed when it merges.

Assisted-by: ClaudeCode:claude-fable-5-1

Author: d-v-b

changes/3285.feature.md

@@ -1,7 +1,7 @@
 JSON metadata validation now delegates to ``msgspec.convert`` for the type
 coercions it supports (``Literal`` membership, ``int`` / ``bool`` strictness,
 list-to-tuple), replacing the per-field hand-written ``parse_*`` logic.
-User-defined attributes retain their existing JSON handling. 
+User-defined attributes retain their existing JSON handling.
 A latent generator-exhaustion bug in
 ``parse_storage_transformers`` is also fixed. See #3285.
 

changes/4331.misc.md

@@ -0,0 +1 @@
+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.

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,
@@ -255,6 +256,30 @@ 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, ...],
+    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: 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,
+        index_codecs=[BytesCodec(), Crc32cCodec()],
+        chunk_shape=chunk_shape,
+    )
+
+
 @st.composite
 def np_array_and_chunks(
     draw: st.DrawFn,
@@ -334,14 +359,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 +551,83 @@ 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,
+    nested: bool | None = None,
+) -> 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 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")
+    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=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 == 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
+
+
 def is_negative_slice(idx: Any) -> bool:
     return isinstance(idx, slice) and idx.step is not None and idx.step < 0
 

tests/test_indexing.py

@@ -2234,14 +2234,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)

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[:]
@@ -182,6 +190,14 @@ async def test_oindex(data: st.DataObject) -> None:
             # behaviour of setitem with repeated indices is not guaranteed in practice
             # Negative and positive spellings of the same index are duplicates too.
             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
@@ -198,6 +214,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[:]