test: give the sharded coordinate-selection write path real weight in test_oindex

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

@@ -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)

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[:]