fix(indexing): read a Boolean array index the way NumPy does

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.

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:

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)

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