From fd4ac097c9f3bbade1dcacbd04e2b133e3ebc12f Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 6 Nov 2026 22:05:54 +0100 Subject: [PATCH 1/3] fix(indexing): read a Boolean array index the way NumPy does A Boolean array used as the whole index now selects what NumPy selects. A 1-D mask over the first of several axes is read as an orthogonal selection, an N-D mask over the leading axes as coordinates broadcast against the trailing axes, and an empty mask selects nothing whatever its shape. Full-shape masks keep the mask selection path, including on arrays with a zero-length axis. Orthogonal selection also accepts an empty Boolean array, and no longer fails on a zero-length dimension. Assisted-by: ClaudeCode:claude-opus-5-5 --- src/zarr/core/array.py | 11 ++++--- src/zarr/core/indexing.py | 67 +++++++++++++++++++++++++++++++++++++-- tests/test_indexing.py | 65 +++++++++++++++++++++++++++++++++++++ 3 files changed, 136 insertions(+), 7 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 8c473f3223..03e91f15dc 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -106,6 +106,7 @@ _iter_grid, check_fields, check_no_multi_fields, + expand_bool_index, is_pure_fancy_indexing, is_pure_orthogonal_indexing, is_scalar, @@ -2758,9 +2759,10 @@ def __getitem__(self, selection: Selection) -> NDArrayLikeOrScalar: [vindex][zarr.Array.vindex], [oindex][zarr.Array.oindex], [blocks][zarr.Array.blocks], [__setitem__][zarr.Array.__setitem__] """ - fields, pure_selection = pop_fields(selection) + expanded = expand_bool_index(selection, self.shape) + fields, pure_selection = pop_fields(expanded) if is_pure_fancy_indexing(pure_selection, self.ndim): - return self.vindex[cast("CoordinateSelection | MaskSelection", selection)] + return self.vindex[cast("CoordinateSelection | MaskSelection", expanded)] elif is_pure_orthogonal_indexing(pure_selection, self.ndim): return self.get_orthogonal_selection(pure_selection, fields=fields) else: @@ -2871,9 +2873,10 @@ def __setitem__(self, selection: Selection, value: npt.ArrayLike) -> None: # which reads the data synchronously before we enter the async context. if isinstance(value, Array): value = np.asarray(value) - fields, pure_selection = pop_fields(selection) + expanded = expand_bool_index(selection, self.shape) + fields, pure_selection = pop_fields(expanded) if is_pure_fancy_indexing(pure_selection, self.ndim): - self.vindex[cast("CoordinateSelection | MaskSelection", selection)] = value + self.vindex[cast("CoordinateSelection | MaskSelection", expanded)] = value elif is_pure_orthogonal_indexing(pure_selection, self.ndim): self.set_orthogonal_selection(pure_selection, value, fields=fields) else: diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index 8ea59eb960..3fcd71f349 100644 --- a/src/zarr/core/indexing.py +++ b/src/zarr/core/indexing.py @@ -642,8 +642,8 @@ def __init__( if not is_bool_array(dim_sel, 1): raise IndexError("Boolean arrays in an orthogonal selection must be 1-dimensional only") - # check shape - if dim_sel.shape[0] != dim_len: + # check shape; an empty Boolean array selects nothing, as in NumPy + if dim_sel.shape[0] not in (dim_len, 0): raise IndexError( f"Boolean array has the wrong length for dimension; expected {dim_len}, got {dim_sel.shape[0]}" ) @@ -660,7 +660,8 @@ def __init__( dim_sel[dim_offset : dim_offset + chunk_len] ) chunk_nitems_cumsum = np.cumsum(chunk_nitems) - nitems = chunk_nitems_cumsum[-1] + # a zero-length dimension has no chunks, so the cumulative sum may be empty + nitems = int(chunk_nitems.sum()) dim_chunk_ixs = np.nonzero(chunk_nitems)[0] # store attributes @@ -1566,6 +1567,65 @@ def pop_fields(selection: SelectionWithFields) -> tuple[Fields | None, Selection return fields, selection +def expand_bool_index( + selection: SelectionWithFields, shape: tuple[int, ...] +) -> SelectionWithFields: + """Rewrite a Boolean array used as the whole index into an equivalent selection. + + NumPy reads such an array `m` as `m.nonzero()`: coordinates along the leading + `m.ndim` axes, with the remaining axes taken whole. A mask with the full shape of + the array stays a mask selection. A 1-dimensional mask over the first of several + axes becomes an orthogonal selection. Any other mask, including an empty one + whose shape does not match (which NumPy accepts as selecting nothing), becomes a + coordinate selection that broadcasts against the trailing axes. All other + selections are returned unchanged. + + Raises + ------ + IndexError + If a non-empty mask does not match the shape of the leading axes it covers. + """ + fields, pure_selection = pop_fields(selection) + if ( + not is_bool_array(pure_selection) + or not 0 < pure_selection.ndim <= len(shape) + or pure_selection.shape == shape + ): + return selection + + mask = np.asarray(pure_selection) + # as in NumPy, an empty mask may have length 0 along any axis it covers + lengths_allowed = {0} if mask.size == 0 else set() + for axis, (mask_len, dim_len) in enumerate(zip(mask.shape, shape, strict=False)): + if mask_len not in lengths_allowed | {dim_len}: + raise IndexError( + f"Boolean index did not match indexed array along axis {axis}; size of " + f"axis is {dim_len} but size of corresponding Boolean axis is {mask_len}" + ) + + expanded: tuple[Any, ...] + if mask.ndim == 1 and mask.size > 0: + # spell out the trailing slices so the selection reads as orthogonal + expanded = (mask,) + (slice(None),) * (len(shape) - 1) + else: + # The trailing axes are read in full, so this materializes one coordinate + # per selected element. + n_trailing = len(shape) - mask.ndim + leading = tuple(c.reshape(c.shape + (1,) * n_trailing) for c in np.nonzero(mask)) + trailing = tuple( + np.arange(dim_len, dtype=np.intp).reshape( + (1,) * (axis + 1) + (dim_len,) + (1,) * (n_trailing - axis - 1) + ) + for axis, dim_len in enumerate(shape[mask.ndim :]) + ) + expanded = leading + trailing + + if fields is not None: + expanded += (fields,) if isinstance(fields, str) else tuple(fields) + # fields mixed into a selection tuple are what `pop_fields` takes apart + return cast("SelectionWithFields", expanded) + + def make_slice_selection(selection: Any) -> list[slice]: ls: list[slice] = [] for dim_selection in selection: @@ -1746,6 +1806,7 @@ def colexicographic_order_coords(shape: tuple[int, ...]) -> tuple[tuple[int, ... def get_indexer( selection: SelectionWithFields, shape: tuple[int, ...], chunk_grid: ChunkGrid ) -> Indexer: + selection = expand_bool_index(selection, shape) _, pure_selection = pop_fields(selection) if is_pure_fancy_indexing(pure_selection, len(shape)): new_selection = ensure_tuple(selection) diff --git a/tests/test_indexing.py b/tests/test_indexing.py index 785f089513..333519a007 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -428,6 +428,71 @@ def test_get_basic_selection_2d_rejects_list_in_tuple(store: StorePath) -> None: z.get_basic_selection((slice(None), [0, 1])) +_EVEN_ROWS = np.array([True, False, True, False]) +_EMPTY_MASK = np.array([], dtype=bool) + + +@pytest.mark.parametrize( + ("shape", "selection"), + [ + # full-shape masks, including ones with a zero-length axis + ((4,), _EVEN_ROWS), + ((4, 3), np.eye(4, 3, dtype=bool)), + ((0,), _EMPTY_MASK), + ((0, 3), np.zeros((0, 3), dtype=bool)), + ((4, 0), np.zeros((4, 0), dtype=bool)), + # a 1-D mask over the first of several axes + ((4, 3), _EVEN_ROWS), + ((4, 3), np.ones(4, dtype=bool)), + ((4, 3), np.zeros(4, dtype=bool)), + ((4, 0), _EVEN_ROWS), + # an N-D mask over the leading axes of a higher-dimensional array + ((2, 3, 4), np.array([[True, False, True], [False, True, True]])), + ((2, 3, 4, 5), np.array([[True, False, True], [False, True, True]])), + # an empty mask selects nothing, whatever axes it covers + ((4,), _EMPTY_MASK), + ((4, 3), _EMPTY_MASK), + ((0, 3), _EMPTY_MASK), + ((2, 3, 4), np.zeros((0, 3), dtype=bool)), + # empty masks inside a selection tuple + ((4, 3), (_EMPTY_MASK, slice(None))), + ((4, 3), (slice(None), _EMPTY_MASK)), + ((0, 3), (_EMPTY_MASK, slice(None))), + ], +) +def test_bool_index_matches_numpy(store: StorePath, shape: tuple[int, ...], selection: Any) -> None: + """A Boolean array index selects, and writes, the same elements as in NumPy.""" + a = np.arange(np.prod(shape), dtype="i4").reshape(shape) + z = zarr_array_from_numpy_array(store, a, chunk_shape=tuple(max(1, s // 2) for s in shape)) + assert_array_equal(z[selection], a[selection]) + + z[selection] = -1 + a[selection] = -1 + assert_array_equal(z[...], a) + + +@pytest.mark.parametrize( + ("shape", "selection"), + [ + ((4,), np.array([True, False, True])), + ((4, 3), np.array([True, False, True])), + ((2, 3, 4), np.ones((2, 2), dtype=bool)), + # an empty mask still has to match along every axis it does not leave empty + ((5, 2), np.zeros((2, 0), dtype=bool)), + ((2, 3, 4), np.zeros((7, 0), dtype=bool)), + ], +) +def test_bool_index_wrong_shape_raises( + store: StorePath, shape: tuple[int, ...], selection: npt.NDArray[np.bool_] +) -> None: + """A non-empty Boolean array index must match the leading axes it covers.""" + z = zarr_array_from_numpy_array(store, np.zeros(shape, dtype="i4")) + with pytest.raises(IndexError, match="Boolean index did not match"): + z[selection] + with pytest.raises(IndexError, match="Boolean index did not match"): + z[selection] = 1 + + def test_fancy_indexing_fallback_on_get_setitem(store: StorePath) -> None: """Paired integer-list indexing falls back to vectorized (fancy) get and set via `__getitem__`/`__setitem__`.""" z = zarr_array_from_numpy_array(store, np.zeros((20, 20))) From 3549811637d9967bc888a2ad61dbd31e5ce2b550 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 6 Nov 2026 22:11:44 +0100 Subject: [PATCH 2/3] docs: add the changelog entry for the Boolean index fix Assisted-by: ClaudeCode:claude-opus-5-5 --- changes/443.bugfix.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changes/443.bugfix.md diff --git a/changes/443.bugfix.md b/changes/443.bugfix.md new file mode 100644 index 0000000000..eaa4f97cc8 --- /dev/null +++ b/changes/443.bugfix.md @@ -0,0 +1 @@ +Indexing an array with a Boolean NumPy array now selects what NumPy selects. That covers a mask over the leading axes (for example a `(2, 3)` mask on a `(2, 3, 4)` array), an empty mask whose length is 0 along an axis, and a full-shape mask on an array with a zero-length axis. Orthogonal selection also accepts an empty Boolean array and no longer fails on a zero-length dimension. From 9cdcca7dcd67b4b8229eef5317a1cb86669c3c87 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 6 Nov 2026 23:26:04 +0100 Subject: [PATCH 3/3] fix(indexing): explain that an N-D Boolean array must be the whole index `z[m2, :]` with a 2-D mask used to fail with "Boolean arrays in an orthogonal selection must be 1-dimensional only", which names an indexing mode the caller never asked for. The error now says that such an array is only supported as the whole index and suggests alternatives. Assisted-by: ClaudeCode:claude-opus-5-5 --- changes/443.bugfix.md | 2 +- src/zarr/core/indexing.py | 13 ++++++++++++- tests/test_indexing.py | 18 ++++++++++++++++++ 3 files changed, 31 insertions(+), 2 deletions(-) diff --git a/changes/443.bugfix.md b/changes/443.bugfix.md index eaa4f97cc8..293bcbd284 100644 --- a/changes/443.bugfix.md +++ b/changes/443.bugfix.md @@ -1 +1 @@ -Indexing an array with a Boolean NumPy array now selects what NumPy selects. That covers a mask over the leading axes (for example a `(2, 3)` mask on a `(2, 3, 4)` array), an empty mask whose length is 0 along an axis, and a full-shape mask on an array with a zero-length axis. Orthogonal selection also accepts an empty Boolean array and no longer fails on a zero-length dimension. +Indexing an array with a Boolean NumPy array now selects what NumPy selects. That covers a mask over the leading axes (for example a `(2, 3)` mask on a `(2, 3, 4)` array), an empty mask whose length is 0 along an axis, and a full-shape mask on an array with a zero-length axis. Orthogonal selection also accepts an empty Boolean array and no longer fails on a zero-length dimension. A Boolean array with more than one dimension inside a selection tuple, which is not supported, now raises an error that says so instead of one about orthogonal selection. diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index 3fcd71f349..4a8dd363d8 100644 --- a/src/zarr/core/indexing.py +++ b/src/zarr/core/indexing.py @@ -1583,9 +1583,20 @@ def expand_bool_index( Raises ------ IndexError - If a non-empty mask does not match the shape of the leading axes it covers. + If a non-empty mask does not match the shape of the leading axes it covers, + or if a Boolean array with more than one dimension appears inside a tuple. """ fields, pure_selection = pop_fields(selection) + if isinstance(pure_selection, tuple): + for position, item in enumerate(pure_selection): + if is_bool_array(item) and len(item.shape) > 1: + raise IndexError( + f"got a {len(item.shape)}-dimensional Boolean array at position {position} " + "of the selection; a Boolean array with more than one dimension is only " + "supported as the whole index, as in `array[mask]`. Index with the mask " + "alone, or read the region first and index the result with NumPy." + ) + return selection if ( not is_bool_array(pure_selection) or not 0 < pure_selection.ndim <= len(shape) diff --git a/tests/test_indexing.py b/tests/test_indexing.py index 333519a007..ed9bd4956a 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -493,6 +493,24 @@ def test_bool_index_wrong_shape_raises( z[selection] = 1 +@pytest.mark.parametrize( + "selection", + [ + (np.ones((2, 3), dtype=bool), slice(None)), + (np.ones((2, 3), dtype=bool), 1), + (Ellipsis, np.ones((3, 4), dtype=bool)), + (slice(None), np.ones((3, 4), dtype=bool)), + ], +) +def test_bool_index_nd_in_tuple_raises(store: StorePath, selection: tuple[Any, ...]) -> None: + """A Boolean array with more than one dimension is only supported as the whole index.""" + z = zarr_array_from_numpy_array(store, np.zeros((2, 3, 4), dtype="i4")) + with pytest.raises(IndexError, match="only supported as the whole index"): + z[selection] + with pytest.raises(IndexError, match="only supported as the whole index"): + z[selection] = 1 + + def test_fancy_indexing_fallback_on_get_setitem(store: StorePath) -> None: """Paired integer-list indexing falls back to vectorized (fancy) get and set via `__getitem__`/`__setitem__`.""" z = zarr_array_from_numpy_array(store, np.zeros((20, 20)))