(feat): support ellipsis indexing

src/anndata/_core/index.py

@@ -3,7 +3,6 @@
 from collections.abc import Iterable, Sequence
 from functools import singledispatch
 from itertools import repeat
-from types import EllipsisType
 from typing import TYPE_CHECKING
 
 import h5py
@@ -14,7 +13,7 @@
 from ..compat import AwkArray, DaskArray, SpArray
 
 if TYPE_CHECKING:
-    from ..compat import Index, IndexRest
+    from ..compat import Index, Index1D
 
 
 def _normalize_indices(
@@ -30,12 +29,10 @@
         if len(index) > 2:
             # only one remaining ellipsis
             num_ellipsis = sum(i is Ellipsis for i in index)
-            if num_ellipsis == 1:
-                index = tuple(i for i in index if i is not Ellipsis)
-            elif num_ellipsis > 1:
+            if num_ellipsis > 1:
                 raise IndexError("an index can only have a single ellipsis ('...')")
-            if len(index) > 2:
+            if num_ellipsis == 0:
                 raise ValueError("AnnData can only be sliced in rows and columns.")
         # deal with pd.Series
         # TODO: The series should probably be aligned first
         if isinstance(index[1], pd.Series):
@@ -55,8 +52,7 @@
     | str
     | Sequence[bool | int | np.integer]
     | np.ndarray
-    | pd.Index
-    | EllipsisType,
+    | pd.Index,
     index: pd.Index,
 ) -> slice | int | np.ndarray:  # ndarray of int or bool
     if not isinstance(index, pd.RangeIndex):
@@ -81,8 +77,6 @@
         return slice(start, stop, step)
     elif isinstance(indexer, np.integer | int):
         return indexer
-    elif isinstance(indexer, EllipsisType):
-        return slice(None)
     elif isinstance(indexer, str):
         return index.get_loc(indexer)  # int
     elif isinstance(
@@ -118,7 +112,7 @@
                     "are not valid obs/ var names or indices."
                 )
             return positions  # np.ndarray[int]
     raise IndexError(f"Unknown indexer {indexer!r} of type {type(indexer)}")
 
 
 def _fix_slice_bounds(s: slice, length: int) -> slice:
@@ -140,15 +134,30 @@
     return slice(start, stop, step)
 
 
-def unpack_index(index: Index) -> tuple[IndexRest, IndexRest]:
+def unpack_index(index: Index) -> tuple[Index1D, Index1D]:
     if not isinstance(index, tuple):
+        if index is Ellipsis:
+            index = slice(None)
         return index, slice(None)
-    elif len(index) == 2:
+    num_ellipsis = sum(i is Ellipsis for i in index)
+    if num_ellipsis > 1:
+        raise IndexError("an index can only have a single ellipsis ('...')")
+    # If index has Ellipsis, filter it out (and if not, error)
+    if len(index) > 2:
+        if not num_ellipsis:
+            raise IndexError("Received a length 3 index without an ellipsis")
+        index = tuple(i for i in index if i is not Ellipsis)
         return index
-    elif len(index) == 1:
-        return index[0], slice(None)
-    else:
-        raise IndexError("invalid number of indices")
+    # If index has Ellipsis, replace it with slice
+    if len(index) == 2:
+        index = tuple(slice(None) if i is Ellipsis else i for i in index)
+        return index
+    if len(index) == 1:
+        index = index[0]
+        if index is Ellipsis:
+            index = slice(None)
+        return index, slice(None)
+    raise IndexError("invalid number of indices")
 
 
 @singledispatch

src/anndata/compat/__init__.py

@@ -53,6 +53,8 @@ class Empty:
     IndexRest
     | tuple[Index1D, IndexRest]
     | tuple[IndexRest, Index1D]
+    | tuple[Index1D, Index1D, EllipsisType]
+    | tuple[EllipsisType, Index1D, Index1D]
     | scipy.sparse.spmatrix
     | SpArray
 )

tests/test_backed_sparse.py

@@ -127,6 +127,28 @@ def test_backed_indexing(
     assert_equal(csr_mem[:, var_idx].X, dense_disk[:, var_idx].X)
 
 
+@pytest.mark.parametrize(
+    "indexing_func",
+    [
+        lambda x: (Ellipsis, slice(0, 10)),
+        lambda x: (slice(0, 10), Ellipsis),
+        lambda x: (slice(0, 10), slice(0, 10), Ellipsis),
+        lambda x: (Ellipsis, slice(0, 10), slice(0, 10)),
+    ],
+    ids=["obs-ellipsis", "var-ellipsis", "obs-var-ellipsis", "ellipsis-obs-var"],
+)
+def test_backed_ellipsis_indexing(
+    ondisk_equivalent_adata: tuple[AnnData, AnnData, AnnData, AnnData],
+    indexing_func,
+):
+    csr_mem, csr_disk, csc_disk, _ = ondisk_equivalent_adata
+
+    index = indexing_func(csr_mem)
+
+    assert_equal(csr_mem.X[index], csr_disk.X[index])
+    assert_equal(csr_disk.X[index], csc_disk.X[index])
+
+
 def make_randomized_mask(size: int) -> np.ndarray:
     randomized_mask = np.zeros(size, dtype=bool)
     inds = np.random.choice(size, 20, replace=False)