diff --git a/changes/4331.misc.md b/changes/4331.misc.md new file mode 100644 index 0000000000..949ebfffd3 --- /dev/null +++ b/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. diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index dfdf25eafd..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, @@ -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 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[:]