(feat): read_lazy for whole AnnData lazy-loading + xarray reading + read_elem_as_dask -> read_elem_lazy

.azure-pipelines.yml

@@ -53,7 +53,7 @@ jobs:
       - script: printf "llvmlite>=0.43\nscanpy>=1.10.0rc1" | tee /tmp/constraints.txt
         displayName: "Create constraints file for `pre-release` and `latest` jobs"
 
-      - script: uv pip install --system --compile "anndata[dev,test] @ ." -c /tmp/constraints.txt
+      - script: uv pip install --system --compile "anndata[dev,test-full] @ ." -c /tmp/constraints.txt
         displayName: "Install dependencies"
         condition: eq(variables['DEPENDENCIES_VERSION'], 'latest')
 
@@ -65,7 +65,7 @@ jobs:
         displayName: "Install minimum dependencies"
         condition: eq(variables['DEPENDENCIES_VERSION'], 'minimum')
 
-      - script: uv pip install -v --system --compile --pre "anndata[dev,test] @ ." -c /tmp/constraints.txt
+      - script: uv pip install -v --system --compile --pre "anndata[dev,test-full] @ ." -c /tmp/constraints.txt
         displayName: "Install dependencies release candidates"
         condition: eq(variables['DEPENDENCIES_VERSION'], 'pre-release')
 
@@ -76,6 +76,10 @@ jobs:
         displayName: "PyTest"
         condition: eq(variables['TEST_TYPE'], 'standard')
 
+      - script: pytest
+        displayName: "PyTest (minimum)"
+        condition: eq(variables['DEPENDENCIES_VERSION'], 'minimum')
+
       - script: pytest --cov --cov-report=xml --cov-context=test
         displayName: "PyTest (coverage)"
         condition: eq(variables['TEST_TYPE'], 'coverage')

.gitignore

@@ -7,6 +7,7 @@ __pycache__/
 /*cache/
 /node_modules/
 /data/
+/venv/
 
 # Distribution / packaging
 /dist/

docs/api.md

@@ -131,7 +131,8 @@ Low level methods for reading and writing elements of an {class}`AnnData` object
 .. autosummary::
    :toctree: generated/
 
-   experimental.read_elem_as_dask
+   experimental.read_elem_lazy
+   experimental.read_lazy
 ```
 
 Utilities for customizing the IO process:
@@ -156,6 +157,9 @@ Types used by the former:
    experimental.ReadCallback
    experimental.WriteCallback
    experimental.StorageType
+   experimental.backed._lazy_arrays.MaskedArray
+   experimental.backed._lazy_arrays.CategoricalArray
+   experimental.backed._xarray.Dataset2D
 ```
 
 ## Errors and warnings

docs/conf.py

@@ -108,6 +108,7 @@
     ("py:class", "numpy.ma.core.MaskedArray"),
     ("py:class", "dask.array.core.Array"),
     ("py:class", "awkward.highlevel.Array"),
+    ("py:class", "awkward.Array"),
     ("py:class", "anndata._core.sparse_dataset.BaseCompressedSparseDataset"),
     ("py:obj", "numpy._typing._array_like._ScalarType_co"),
     # https://github.com/sphinx-doc/sphinx/issues/10974

docs/release-notes/0.11.0rc1.md

@@ -20,8 +20,8 @@
 - `scipy.sparse.csr_array` and `scipy.sparse.csc_array` are now supported when constructing `AnnData` objects {user}`ilan-gold` {user}`isaac-virshup` ({pr}`1028`)
 - Allow `axis` parameter of e.g. :func:`anndata.concat` to accept `'obs'` and `'var'` {user}`flying-sheep` ({pr}`1244`)
 - Add `settings` object with methods for altering internally-used options, like checking for uniqueness on `obs`' index {user}`ilan-gold` ({pr}`1270`)
+- Add :func:`~anndata.experimental.read_elem_lazy` function to handle i/o with sparse and dense arrays {user}`ilan-gold` ({pr}`1469`)
 - Add {attr}`~anndata.settings.remove_unused_categories` option to {attr}`anndata.settings` to override current behavior {user}`ilan-gold` ({pr}`1340`)
-- Add :func:`~anndata.experimental.read_elem_as_dask` function to handle i/o with sparse and dense arrays {user}`ilan-gold` ({pr}`1469`)
 - Add ability to convert strings to categoricals on write in {meth}`~anndata.AnnData.write_h5ad` and {meth}`~anndata.AnnData.write_zarr` via `convert_strings_to_categoricals` parameter {user}` falexwolf` ({pr}`1474`)
 - Add {attr}`~anndata.settings.check_uniqueness` option to {attr}`anndata.settings` to override current behavior {user}`ilan-gold` ({pr}`1507`)
 - Add functionality to write from GPU {class}`dask.array.Array` to disk {user}`ilan-gold` ({pr}`1550`)

docs/release-notes/1247.feature.md

@@ -0,0 +1 @@
+Add {func}`~anndata.experimental.read_elem_lazy` (in place of `read_elem_as_dask`) to handle backed dataframes, sparse arrays, and dense arrays, as well as a {func}`~anndata.experimental.read_lazy` to handle reading in as much of the on-disk data as possible to produce a {class}`~anndata.AnnData` object {user}`ilan-gold`

docs/tutorials/index.md

@@ -14,4 +14,5 @@ notebooks/anncollection-annloader
 notebooks/anndata_dask_array
 notebooks/awkward-arrays
 notebooks/{read,write}_dispatched
+notebooks/read_lazy
 ```

pyproject.toml

@@ -84,6 +84,7 @@ doc = [
     "anndata[dev-doc]",
 ]
 dev-doc = ["towncrier>=24.8.0"] # release notes tool
+test-full = ["anndata[test]", "anndata[lazy]"]
 test = [
     "loompy>=3.0.5",
     "pytest>=8.2",
@@ -109,6 +110,7 @@ cu12 = ["cupy-cuda12x"]
 cu11 = ["cupy-cuda11x"]
 # https://github.com/dask/dask/issues/11290
 dask = ["dask[array]>=2022.09.2,<2024.8.0"]
+lazy = ["xarray>=2024.06.0", "aiohttp", "requests", "zarr<3.0.0a0", "anndata[dask]"]
 
 [tool.hatch.version]
 source = "vcs"

src/anndata/_core/aligned_df.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import warnings
+from collections.abc import Mapping
 from functools import singledispatch
 from typing import TYPE_CHECKING
 
@@ -10,13 +11,26 @@
 from .._warnings import ImplicitModificationWarning
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable, Mapping
+    from collections.abc import Iterable
     from typing import Any, Literal
 
 
 @singledispatch
 def _gen_dataframe(
-    anno: Mapping[str, Any],
+    anno: Any,
+    index_names: Iterable[str],
+    *,
+    source: Literal["X", "shape"],
+    attr: Literal["obs", "var"],
+    length: int | None = None,
+) -> pd.DataFrame:
+    raise ValueError(f"Cannot convert {type(anno)} to {attr} DataFrame")
+
+
+@_gen_dataframe.register(Mapping)
+@_gen_dataframe.register(type(None))
+def _gen_dataframe_mapping(
+    anno: Mapping[str, Any] | None,
     index_names: Iterable[str],
     *,
     source: Literal["X", "shape"],

src/anndata/_core/anndata.py

@@ -8,7 +8,7 @@
 from collections import OrderedDict
 from collections.abc import Mapping, MutableMapping, Sequence
 from copy import copy, deepcopy
-from functools import partial
+from functools import partial, singledispatch
 from pathlib import Path
 from textwrap import dedent
 from typing import TYPE_CHECKING
@@ -41,7 +41,6 @@
 from .sparse_dataset import BaseCompressedSparseDataset, sparse_dataset
 from .storage import coerce_array
 from .views import (
-    DataFrameView,
     DictView,
     _resolve_idxs,
     as_view,
@@ -301,8 +300,8 @@ def _init_as_view(self, adata_ref: AnnData, oidx: Index, vidx: Index):
             self._remove_unused_categories(adata_ref.obs, obs_sub, uns)
             self._remove_unused_categories(adata_ref.var, var_sub, uns)
         # set attributes
-        self._obs = DataFrameView(obs_sub, view_args=(self, "obs"))
-        self._var = DataFrameView(var_sub, view_args=(self, "var"))
+        self._obs = as_view(obs_sub, view_args=(self, "obs"))
+        self._var = as_view(var_sub, view_args=(self, "var"))
         self._uns = uns
 
         # set data
@@ -1023,8 +1022,10 @@ def __getitem__(self, index: Index) -> AnnData:
         oidx, vidx = self._normalize_indices(index)
         return AnnData(self, oidx=oidx, vidx=vidx, asview=True)
 
+    @staticmethod
+    @singledispatch
     def _remove_unused_categories(
-        self, df_full: pd.DataFrame, df_sub: pd.DataFrame, uns: dict[str, Any]
+        df_full: pd.DataFrame, df_sub: pd.DataFrame, uns: dict[str, Any]
     ):
         for k in df_full:
             if not isinstance(df_full[k].dtype, pd.CategoricalDtype):

src/anndata/_core/index.py

@@ -50,10 +50,13 @@
     | pd.Index,
     index: pd.Index,
 ) -> slice | int | np.ndarray:  # ndarray of int or bool
-    if not isinstance(index, pd.RangeIndex):
-        msg = "Don’t call _normalize_index with non-categorical/string names"
-        assert index.dtype != float, msg
-        assert index.dtype != int, msg
+    from ..experimental.backed._compat import xr
+
+    # TODO: why is this here? All tests pass without it and it seems at the minimum not strict enough.
+    # if not isinstance(index, pd.RangeIndex):
+    #     msg = "Don’t call _normalize_index with non-categorical/string names and non-range index"
+    #     assert index.dtype != float, msg
+    #     assert index.dtype != int, msg
 
     # the following is insanely slow for sequences,
     # we replaced it using pandas below
@@ -107,6 +110,10 @@
                     "are not valid obs/ var names or indices."
                 )
             return positions  # np.ndarray[int]
+    elif isinstance(indexer, xr.DataArray):
+        if isinstance(indexer.data, DaskArray):
+            return indexer.data.compute()
+        return indexer.data
     else:
         raise IndexError(f"Unknown indexer {indexer!r} of type {type(indexer)}")