diff --git a/changes/443.bugfix.md b/changes/443.bugfix.md new file mode 100644 index 0000000000..293bcbd284 --- /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. 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/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..4a8dd363d8 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,76 @@ 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, + 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) + 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 +1817,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..ed9bd4956a 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -428,6 +428,89 @@ 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 + + +@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)))