From c6a27a2eb9a0d6d4f7b3a75d5ca50ee7889b6486 Mon Sep 17 00:00:00 2001 From: Max Jones <14077947+maxrjones@users.noreply.github.com> Date: Mon, 8 Jun 2026 16:02:27 -0400 Subject: [PATCH 1/3] test: add property tests for block and mask indexing --- changes/4054.misc.md | 1 + src/zarr/testing/strategies.py | 45 ++++++++++++++++++++++++++ tests/test_properties.py | 59 ++++++++++++++++++++++++++++++++++ 3 files changed, 105 insertions(+) create mode 100644 changes/4054.misc.md diff --git a/changes/4054.misc.md b/changes/4054.misc.md new file mode 100644 index 0000000000..4aa9435b84 --- /dev/null +++ b/changes/4054.misc.md @@ -0,0 +1 @@ +Add Hypothesis property tests for block and mask indexing (`test_block_indexing`, `test_mask_indexing`), along with a `block_indices` strategy in `zarr.testing.strategies`. These extend the existing randomized indexing coverage (basic, orthogonal, and vectorized) to the block and mask selection methods. diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index 74d9d7c683..8699cda5e6 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -591,6 +591,51 @@ def orthogonal_indices( return tuple(zindexer), tuple(np.broadcast_arrays(*npindexer)) +@st.composite +def block_indices( + draw: st.DrawFn, *, shape: tuple[int, ...], chunks: tuple[int, ...] +) -> tuple[tuple[int | slice, ...], tuple[slice, ...]]: + """ + Strategy for block-selection indexers over a *regular* chunk grid. + + Block indexing addresses whole chunks on the block grid rather than + individual elements. It only supports integers and step-1 slices over the + grid (strided block slices are rejected), so neither newaxis, ellipsis, nor + a step is generated here. The array-space translation below assumes a + regular (uniform) chunk grid, so ``shape`` must be evenly tiled by + ``chunks`` up to a possibly-smaller last chunk per dimension. Every + dimension must have at least one chunk (``size >= 1``). + + Returns + ------- + block_indexer + A tuple of ints / step-1 slices addressing whole chunks, suitable for + ``Array.blocks`` / ``Array.get_block_selection`` / ``set_block_selection``. + array_indexer + The equivalent array-space selection (a tuple of slices) for indexing + the corresponding numpy array, used as the comparison oracle. + """ + grid_shape = tuple(-(-s // c) for s, c in zip(shape, chunks, strict=True)) # ceil division + block_indexer: list[int | slice] = [] + array_indexer: list[slice] = [] + for size, chunk, nchunks in zip(shape, chunks, grid_shape, strict=True): + if draw(st.booleans()): + # a single block, sometimes addressed from the end with a negative index + block = draw(st.integers(min_value=-nchunks, max_value=nchunks - 1)) + block_indexer.append(block) + start = (block % nchunks) * chunk + array_indexer.append(slice(start, min(start + chunk, size))) + else: + # a contiguous run of whole blocks (possibly empty). The start must + # reference an existing chunk: block indexing rejects a slice that + # starts at nchunks, unlike numpy which treats arr[len:len] as empty. + start_block = draw(st.integers(min_value=0, max_value=nchunks - 1)) + stop_block = draw(st.integers(min_value=start_block, max_value=nchunks)) + block_indexer.append(slice(start_block, stop_block)) + array_indexer.append(slice(start_block * chunk, min(stop_block * chunk, size))) + return tuple(block_indexer), tuple(array_indexer) + + def key_ranges( keys: SearchStrategy[str] = node_names, max_size: int = sys.maxsize ) -> SearchStrategy[list[tuple[str, RangeByteRequest]]]: diff --git a/tests/test_properties.py b/tests/test_properties.py index 0e5dcf77b0..41105aeb5c 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -25,7 +25,9 @@ array_metadata, arrays, basic_indices, + block_indices, complex_rectilinear_arrays, + np_array_and_chunks, numpy_arrays, orthogonal_indices, rectilinear_arrays, @@ -230,6 +232,63 @@ async def test_vindex(data: st.DataObject) -> None: # note: async vindex setitem not yet implemented +@settings(deadline=None) +@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") +@given(data=st.data()) +def test_mask_indexing(data: st.DataObject) -> None: + zarray = data.draw(st.one_of(simple_arrays(), rectilinear_arrays())) + nparray = zarray[:] + mask = data.draw(npst.arrays(dtype=np.bool_, shape=st.just(nparray.shape))) + + expected = nparray[mask] + + # sync get, via both the dedicated method and the vindex interface + assert_array_equal(expected, zarray.get_mask_selection(mask)) + assert_array_equal(expected, zarray.vindex[mask]) + + # sync set, via both interfaces + assume(zarray.shards is None) # GH2834 + new_data = data.draw(numpy_arrays(shapes=st.just(expected.shape), dtype=nparray.dtype)) + nparray[mask] = new_data + zarray.set_mask_selection(mask, new_data) + assert_array_equal(nparray, zarray[:]) + + zarray.vindex[mask] = new_data + assert_array_equal(nparray, zarray[:]) + + +@settings(deadline=None) +@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") +@given(data=st.data()) +def test_block_indexing(data: st.DataObject) -> None: + # Block indexing addresses whole chunks on a regular grid; the array-space + # oracle in block_indices() assumes regular, unsharded chunks, so build the + # array directly from a regular chunking rather than drawing one that might + # be rectilinear or sharded. + nparray, chunks = data.draw( + np_array_and_chunks(arrays=numpy_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1))) + ) + store = data.draw(stores) + zarray = zarr.create_array(store=store, shape=nparray.shape, chunks=chunks, dtype=nparray.dtype) + zarray[...] = nparray + + block_indexer, array_indexer = data.draw(block_indices(shape=nparray.shape, chunks=chunks)) + expected = nparray[array_indexer] + + # sync get, via both the .blocks interface and the dedicated method + assert_array_equal(expected, zarray.blocks[block_indexer]) + assert_array_equal(expected, zarray.get_block_selection(block_indexer)) + + # sync set, via both interfaces + new_data = data.draw(numpy_arrays(shapes=st.just(expected.shape), dtype=nparray.dtype)) + nparray[array_indexer] = new_data + zarray.blocks[block_indexer] = new_data + assert_array_equal(nparray, zarray[:]) + + zarray.set_block_selection(block_indexer, new_data) + assert_array_equal(nparray, zarray[:]) + + @given(store=stores, meta=array_metadata()) # type: ignore[misc] @pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") async def test_roundtrip_array_metadata_from_store( From d6a12e6c5761618d670718a7159911dca1d99a1c Mon Sep 17 00:00:00 2001 From: Max Jones <14077947+maxrjones@users.noreply.github.com> Date: Mon, 8 Jun 2026 19:41:32 -0400 Subject: [PATCH 2/3] test: reuse basic_indices in block_indices strategy --- src/zarr/testing/strategies.py | 61 ++++++++++++++++++++++------------ 1 file changed, 39 insertions(+), 22 deletions(-) diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index 8699cda5e6..11adcccb27 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -598,41 +598,58 @@ def block_indices( """ Strategy for block-selection indexers over a *regular* chunk grid. - Block indexing addresses whole chunks on the block grid rather than - individual elements. It only supports integers and step-1 slices over the - grid (strided block slices are rejected), so neither newaxis, ellipsis, nor - a step is generated here. The array-space translation below assumes a - regular (uniform) chunk grid, so ``shape`` must be evenly tiled by - ``chunks`` up to a possibly-smaller last chunk per dimension. Every - dimension must have at least one chunk (``size >= 1``). + Block indexing is basic indexing applied to the block grid (the grid of + chunks), so each axis is drawn with ``basic_indices`` over that axis's chunk + count -- mirroring how ``orthogonal_indices`` reuses ``basic_indices`` per + axis. Block indexing only supports integers and step-1 slices whose start + references an existing chunk, so strided slices and slices starting at the + grid edge are filtered out. The array-space translation assumes a regular + (uniform) chunk grid, so ``shape`` must be evenly tiled by ``chunks`` up to a + possibly-smaller last chunk per dimension, and every dimension must have at + least one chunk (``size >= 1``). Returns ------- block_indexer - A tuple of ints / step-1 slices addressing whole chunks, suitable for - ``Array.blocks`` / ``Array.get_block_selection`` / ``set_block_selection``. + A per-axis tuple of ints / step-1 slices addressing whole chunks, + suitable for ``Array.blocks`` / ``get_block_selection`` / ``set_block_selection``. array_indexer The equivalent array-space selection (a tuple of slices) for indexing the corresponding numpy array, used as the comparison oracle. """ grid_shape = tuple(-(-s // c) for s, c in zip(shape, chunks, strict=True)) # ceil division + + def supported(nchunks: int) -> Callable[[tuple[Any, ...]], bool]: + # Block indexing only accepts step-1 slices whose start references an + # existing chunk (a slice starting at nchunks raises, unlike numpy). + def predicate(value: tuple[Any, ...]) -> bool: + dim_sel = value[0] + if isinstance(dim_sel, slice): + if dim_sel.step not in (None, 1): + return False + start = dim_sel.start or 0 + return 0 <= (start + nchunks if start < 0 else start) < nchunks + return True + + return predicate + block_indexer: list[int | slice] = [] array_indexer: list[slice] = [] for size, chunk, nchunks in zip(shape, chunks, grid_shape, strict=True): - if draw(st.booleans()): - # a single block, sometimes addressed from the end with a negative index - block = draw(st.integers(min_value=-nchunks, max_value=nchunks - 1)) - block_indexer.append(block) - start = (block % nchunks) * chunk - array_indexer.append(slice(start, min(start + chunk, size))) + (dim_sel,) = draw( + basic_indices(min_dims=1, shape=(nchunks,), allow_ellipsis=False) + # normalize bare ints / slices to a 1-tuple, skip the empty tuple + .map(lambda x: (x,) if not isinstance(x, tuple) else x) + .filter(bool) + .filter(supported(nchunks)) + ) + block_indexer.append(dim_sel) + if isinstance(dim_sel, slice): + start, stop, _ = dim_sel.indices(nchunks) + array_indexer.append(slice(start * chunk, min(stop * chunk, size))) else: - # a contiguous run of whole blocks (possibly empty). The start must - # reference an existing chunk: block indexing rejects a slice that - # starts at nchunks, unlike numpy which treats arr[len:len] as empty. - start_block = draw(st.integers(min_value=0, max_value=nchunks - 1)) - stop_block = draw(st.integers(min_value=start_block, max_value=nchunks)) - block_indexer.append(slice(start_block, stop_block)) - array_indexer.append(slice(start_block * chunk, min(stop_block * chunk, size))) + block = dim_sel % nchunks + array_indexer.append(slice(block * chunk, min((block + 1) * chunk, size))) return tuple(block_indexer), tuple(array_indexer) From c206ed249a3fde23a180b22c849258426e9c28bc Mon Sep 17 00:00:00 2001 From: Max Jones <14077947+maxrjones@users.noreply.github.com> Date: Mon, 8 Jun 2026 19:49:39 -0400 Subject: [PATCH 3/3] test: take chunk_grid_shape as input in block_indices strategy --- src/zarr/testing/strategies.py | 23 +++++++++++------------ tests/test_properties.py | 4 +++- 2 files changed, 14 insertions(+), 13 deletions(-) diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index 11adcccb27..7d6556a359 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -593,20 +593,20 @@ def orthogonal_indices( @st.composite def block_indices( - draw: st.DrawFn, *, shape: tuple[int, ...], chunks: tuple[int, ...] + draw: st.DrawFn, *, chunk_grid_shape: tuple[int, ...], chunks: tuple[int, ...] ) -> tuple[tuple[int | slice, ...], tuple[slice, ...]]: """ Strategy for block-selection indexers over a *regular* chunk grid. Block indexing is basic indexing applied to the block grid (the grid of chunks), so each axis is drawn with ``basic_indices`` over that axis's chunk - count -- mirroring how ``orthogonal_indices`` reuses ``basic_indices`` per - axis. Block indexing only supports integers and step-1 slices whose start - references an existing chunk, so strided slices and slices starting at the - grid edge are filtered out. The array-space translation assumes a regular - (uniform) chunk grid, so ``shape`` must be evenly tiled by ``chunks`` up to a - possibly-smaller last chunk per dimension, and every dimension must have at - least one chunk (``size >= 1``). + count from ``chunk_grid_shape`` (e.g. ``Array.cdata_shape``), mirroring how + ``orthogonal_indices`` reuses ``basic_indices`` per axis. Block indexing only + supports integers and step-1 slices whose start references an existing chunk, + so strided slices and slices starting at the grid edge are filtered out. The + array-space translation assumes a regular (uniform) chunk grid; an over-long + stop into a smaller last chunk is left for numpy to clamp when the oracle is + applied. Returns ------- @@ -617,7 +617,6 @@ def block_indices( The equivalent array-space selection (a tuple of slices) for indexing the corresponding numpy array, used as the comparison oracle. """ - grid_shape = tuple(-(-s // c) for s, c in zip(shape, chunks, strict=True)) # ceil division def supported(nchunks: int) -> Callable[[tuple[Any, ...]], bool]: # Block indexing only accepts step-1 slices whose start references an @@ -635,7 +634,7 @@ def predicate(value: tuple[Any, ...]) -> bool: block_indexer: list[int | slice] = [] array_indexer: list[slice] = [] - for size, chunk, nchunks in zip(shape, chunks, grid_shape, strict=True): + for chunk, nchunks in zip(chunks, chunk_grid_shape, strict=True): (dim_sel,) = draw( basic_indices(min_dims=1, shape=(nchunks,), allow_ellipsis=False) # normalize bare ints / slices to a 1-tuple, skip the empty tuple @@ -646,10 +645,10 @@ def predicate(value: tuple[Any, ...]) -> bool: block_indexer.append(dim_sel) if isinstance(dim_sel, slice): start, stop, _ = dim_sel.indices(nchunks) - array_indexer.append(slice(start * chunk, min(stop * chunk, size))) + array_indexer.append(slice(start * chunk, stop * chunk)) else: block = dim_sel % nchunks - array_indexer.append(slice(block * chunk, min((block + 1) * chunk, size))) + array_indexer.append(slice(block * chunk, (block + 1) * chunk)) return tuple(block_indexer), tuple(array_indexer) diff --git a/tests/test_properties.py b/tests/test_properties.py index 41105aeb5c..994510aca0 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -272,7 +272,9 @@ def test_block_indexing(data: st.DataObject) -> None: zarray = zarr.create_array(store=store, shape=nparray.shape, chunks=chunks, dtype=nparray.dtype) zarray[...] = nparray - block_indexer, array_indexer = data.draw(block_indices(shape=nparray.shape, chunks=chunks)) + block_indexer, array_indexer = data.draw( + block_indices(chunk_grid_shape=zarray.cdata_shape, chunks=chunks) + ) expected = nparray[array_indexer] # sync get, via both the .blocks interface and the dedicated method