(feat): support ellipsis indexing

docs/conf.py

@@ -112,6 +112,8 @@
     ("py:obj", "numpy._typing._array_like._ScalarType_co"),
     # https://github.com/sphinx-doc/sphinx/issues/10974
     ("py:class", "numpy.int64"),
+    # https://github.com/tox-dev/sphinx-autodoc-typehints/issues/498
+    ("py:class", "types.EllipsisType"),
 ]
 
 

docs/release-notes/1729.feature.md

@@ -0,0 +1 @@
+Add support for ellipsis indexing of the {class}`~anndata.AnnData` object {user}`ilan-gold`

src/anndata/_core/anndata.py

@@ -52,9 +52,10 @@
     from os import PathLike
     from typing import Any, Literal
 
+    from ..compat import Index1D
     from ..typing import ArrayDataStructureType
     from .aligned_mapping import AxisArraysView, LayersView, PairwiseArraysView
-    from .index import Index, Index1D
+    from .index import Index
 
 
 # for backwards compat

src/anndata/_core/index.py

@@ -27,7 +27,12 @@
         index: Index = index.values
     if isinstance(index, tuple):
         if len(index) > 2:
-            raise ValueError("AnnData can only be sliced in rows and columns.")
+            # only one remaining ellipsis
+            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 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):
@@ -107,8 +112,7 @@
                     "are not valid obs/ var names or indices."
                 )
             return positions  # np.ndarray[int]
-    else:
-        raise IndexError(f"Unknown indexer {indexer!r} of type {type(indexer)}")
+    raise IndexError(f"Unknown indexer {indexer!r} of type {type(indexer)}")
 
 
 def _fix_slice_bounds(s: slice, length: int) -> slice:
@@ -132,13 +136,28 @@
 
 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

@@ -10,6 +10,7 @@
 from importlib.util import find_spec
 from inspect import Parameter, signature
 from pathlib import Path
+from types import EllipsisType
 from typing import TYPE_CHECKING, TypeVar
 from warnings import warn
 
@@ -47,7 +48,16 @@ class Empty:
 
 
 Index1D = slice | int | str | np.int64 | np.ndarray
-Index = Index1D | tuple[Index1D, Index1D] | scipy.sparse.spmatrix | SpArray
+IndexRest = Index1D | EllipsisType
+Index = (
+    IndexRest
+    | tuple[Index1D, IndexRest]
+    | tuple[IndexRest, Index1D]
+    | tuple[Index1D, Index1D, EllipsisType]
+    | tuple[EllipsisType, Index1D, Index1D]
+    | scipy.sparse.spmatrix
+    | SpArray
+)
 H5Group = h5py.Group
 H5Array = h5py.Dataset
 H5File = h5py.File

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)

tests/test_views.py

@@ -3,6 +3,7 @@
 from contextlib import ExitStack
 from copy import deepcopy
 from operator import mul
+from typing import TYPE_CHECKING
 
 import joblib
 import numpy as np
@@ -35,6 +36,9 @@
 )
 from anndata.utils import asarray
 
+if TYPE_CHECKING:
+    from typing import Literal
+
 IGNORE_SPARSE_EFFICIENCY_WARNING = pytest.mark.filterwarnings(
     "ignore:Changing the sparsity structure:scipy.sparse.SparseEfficiencyWarning"
 )
@@ -786,6 +790,49 @@ def test_dataframe_view_index_setting():
     assert a2.obs.index.values.tolist() == ["a", "b"]
 
 
+@pytest.mark.parametrize("axis", ["obs", "var", None])
+def test_ellipsis_index_2d(
+    adata: ad.AnnData, subset_func, matrix_type, axis: Literal["obs", "var"] | None
+):
+    adata = gen_adata((10, 10), X_type=matrix_type, **GEN_ADATA_DASK_ARGS)
+    if axis is not None:
+        axis_subset = subset_func(getattr(adata, f"{axis}_names"))
+        subset_with_ellipsis = (
+            (axis_subset, Ellipsis) if axis == "obs" else (Ellipsis, axis_subset)
+        )
+        subset_with_slice = (
+            (axis_subset, slice(None)) if axis == "obs" else (slice(None), axis_subset)
+        )
+    else:
+        subset_with_ellipsis = Ellipsis
+        subset_with_slice = (slice(None), slice(None))
+    subset_ellipsis = adata[subset_with_ellipsis]
+    subset = adata[subset_with_slice]
+    assert_equal(subset_ellipsis, subset)
+
+
+@pytest.mark.parametrize("ellipsis_index_loc", [0, 1, 2])
+def test_ellipsis_index_3d(
+    adata: ad.AnnData, subset_func, matrix_type, ellipsis_index_loc
+):
+    adata = gen_adata((10, 10), X_type=matrix_type, **GEN_ADATA_DASK_ARGS)
+    obs_subset = subset_func(getattr(adata, "obs_names"))
+    var_subset = subset_func(getattr(adata, "var_names"))
+    generator = iter((obs_subset, var_subset))
+    subset_with_ellipsis = tuple(
+        Ellipsis if i == ellipsis_index_loc else next(generator) for i in range(3)
+    )
+    subset_with_slice = (obs_subset, var_subset)
+    subset_ellipsis = adata[subset_with_ellipsis]
+    subset = adata[subset_with_slice]
+    assert_equal(subset_ellipsis, subset)
+
+
+def test_ellipsis_index_3d_2_ellipses_errors(adata: ad.AnnData):
+    with pytest.raises(IndexError, match=r"only have a single ellipsis"):
+        gen_adata((10, 10))[..., 0, ...]
+
+
 # @pytest.mark.parametrize("dim", ["obs", "var"])
 # @pytest.mark.parametrize(
 #     ("idx", "pat"),