fix: selection with zarr arrays (#2137)

* fix: selection with zarr arrays

* fixup

* Update src/zarr/core/indexing.py

Co-authored-by: Davis Bennett <[email protected]>

* Apply suggestions from code review

* lint

---------

Co-authored-by: Davis Bennett <[email protected]>

Author: jhamman

src/zarr/core/indexing.py

@@ -35,7 +35,6 @@
 ArrayOfIntOrBool = npt.NDArray[np.intp] | npt.NDArray[np.bool_]
 BasicSelector = int | slice | EllipsisType
 Selector = BasicSelector | ArrayOfIntOrBool
-
 BasicSelection = BasicSelector | tuple[BasicSelector, ...]  # also used for BlockIndex
 CoordinateSelection = IntSequence | tuple[IntSequence, ...]
 MaskSelection = npt.NDArray[np.bool_]
@@ -75,6 +74,15 @@ def err_too_many_indices(selection: Any, shape: ChunkCoords) -> None:
     raise IndexError(f"too many indices for array; expected {len(shape)}, got {len(selection)}")
 
 
+def _zarr_array_to_int_or_bool_array(arr: Array) -> npt.NDArray[np.intp] | npt.NDArray[np.bool_]:
+    if arr.dtype.kind in ("i", "b"):
+        return np.asarray(arr)
+    else:
+        raise IndexError(
+            f"Invalid array dtype: {arr.dtype}. Arrays used as indices must be of integer or boolean type"
+        )
+
+
 @runtime_checkable
 class Indexer(Protocol):
     shape: ChunkCoords
@@ -842,7 +850,14 @@ def __iter__(self) -> Iterator[ChunkProjection]:
 class OIndex:
     array: Array
 
-    def __getitem__(self, selection: OrthogonalSelection) -> NDArrayLike:
+    # TODO: develop Array generic and move zarr.Array[np.intp] | zarr.Array[np.bool_] to ArrayOfIntOrBool
+    def __getitem__(self, selection: OrthogonalSelection | Array) -> NDArrayLike:
+        from zarr.core.array import Array
+
+        # if input is a Zarr array, we materialize it now.
+        if isinstance(selection, Array):
+            selection = _zarr_array_to_int_or_bool_array(selection)
+
         fields, new_selection = pop_fields(selection)
         new_selection = ensure_tuple(new_selection)
         new_selection = replace_lists(new_selection)
@@ -1130,7 +1145,13 @@ def __init__(self, selection: MaskSelection, shape: ChunkCoords, chunk_grid: Chu
 class VIndex:
     array: Array
 
-    def __getitem__(self, selection: CoordinateSelection | MaskSelection) -> NDArrayLike:
+    # TODO: develop Array generic and move zarr.Array[np.intp] | zarr.Array[np.bool_] to ArrayOfIntOrBool
+    def __getitem__(self, selection: CoordinateSelection | MaskSelection | Array) -> NDArrayLike:
+        from zarr.core.array import Array
+
+        # if input is a Zarr array, we materialize it now.
+        if isinstance(selection, Array):
+            selection = _zarr_array_to_int_or_bool_array(selection)
         fields, new_selection = pop_fields(selection)
         new_selection = ensure_tuple(new_selection)
         new_selection = replace_lists(new_selection)

tests/v3/test_indexing.py

@@ -1861,3 +1861,20 @@ def test_orthogonal_bool_indexing_like_numpy_ix(
     # note: in python 3.10 z[*selection] is not valid unpacking syntax
     actual = z[(*selection,)]
     assert_array_equal(expected, actual, err_msg=f"{selection=}")
+
+
+def test_indexing_with_zarr_array(store: StorePath) -> None:
+    # regression test for https://github.com/zarr-developers/zarr-python/issues/2133
+    a = np.arange(10)
+    za = zarr.array(a, chunks=2, store=store, path="a")
+    ix = [False, True, False, True, False, True, False, True, False, True]
+    ii = [0, 2, 4, 5]
+
+    zix = zarr.array(ix, chunks=2, store=store, dtype="bool", path="ix")
+    zii = zarr.array(ii, chunks=2, store=store, dtype="i4", path="ii")
+    assert_array_equal(a[ix], za[zix])
+    assert_array_equal(a[ix], za.oindex[zix])
+    assert_array_equal(a[ix], za.vindex[zix])
+
+    assert_array_equal(a[ii], za[zii])
+    assert_array_equal(a[ii], za.oindex[zii])