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

.azure-pipelines.yml

@@ -49,7 +49,7 @@ jobs:
           path: $(uv_cache_dir)
         displayName: Cache pip packages
 
-      - script: uv pip install --system --compile "anndata[dev,test] @ ." -c ci/constraints.txt
+      - script: uv pip install --system --compile "anndata[dev,test-full] @ ." -c ci/constraints.txt
         displayName: "Install dependencies"
         condition: eq(variables['DEPENDENCIES_VERSION'], 'latest')
 
@@ -61,7 +61,7 @@ jobs:
         displayName: "Install minimum dependencies"
         condition: eq(variables['DEPENDENCIES_VERSION'], 'minimum')
 
-      - script: uv pip install -v --system --compile --pre "anndata[dev,test] @ ." -c ci/constraints.txt
+      - script: uv pip install -v --system --compile --pre "anndata[dev,test-full] @ ." -c ci/constraints.txt
         displayName: "Install dependencies release candidates"
         condition: eq(variables['DEPENDENCIES_VERSION'], 'pre-release')
 
@@ -72,6 +72,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/release-notes/0.11.0.md

@@ -40,7 +40,7 @@ Release candidates:
 - {guilabel}`rc1` Allow `axis` parameter of e.g. {func}`anndata.concat` to accept `'obs'` and `'var'` {user}`flying-sheep` ({pr}`1244`)
 - {guilabel}`rc1` Add `settings` object with methods for altering internally-used options, like checking for uniqueness on `obs`' index {user}`ilan-gold` ({pr}`1270`)
 - {guilabel}`rc1` Add {attr}`~anndata.settings.remove_unused_categories` option to {attr}`anndata.settings` to override current behavior {user}`ilan-gold` ({pr}`1340`)
-- {guilabel}`rc1` Add {func}`~anndata.experimental.read_elem_as_dask` function to handle i/o with sparse and dense arrays {user}`ilan-gold` ({pr}`1469`)
+- {guilabel}`rc1` Add `~anndata.experimental.read_elem_as_dask` function to handle i/o with sparse and dense arrays {user}`ilan-gold` ({pr}`1469`)
 - {guilabel}`rc1` 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`)
 - {guilabel}`rc1` Add {attr}`~anndata.settings.check_uniqueness` option to {attr}`anndata.settings` to override current behavior {user}`ilan-gold` ({pr}`1507`)
 - {guilabel}`rc1` Add functionality to write from GPU {class}`dask.array.Array` to disk {user}`ilan-gold` ({pr}`1550`)

docs/release-notes/0.11.1.md

@@ -4,5 +4,5 @@
 ### Bug fixes
 
 - Remove upper pin on `dask` and exclude versions broken with sparse indexing {user}`ilan-gold` ({pr}`1725`)
-- Fix chunking with -1 in `chunks` argument of {func}`~anndata.experimental.read_elem_as_dask` {user}`ilan-gold` ({pr}`1743`)
+- Fix chunking with -1 in `chunks` argument of `~anndata.experimental.read_elem_as_dask` {user}`ilan-gold` ({pr}`1743`)
 - Fix `cupy<0.13` imports in non-gpu environments {user}`ilan-gold` ({pr}`1754`)

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

@@ -81,9 +81,11 @@ doc = [
     "sphinx_design>=0.5.0",
     "readthedocs-sphinx-search",
     # for unreleased changes
-    "anndata[dev-doc]",
+    "anndata[dev-doc,dask]",
+    "awkward>=2.3"
 ]
 dev-doc = ["towncrier>=24.8.0"] # release notes tool
+test-full = ["anndata[test,lazy]"]
 test = [
     "loompy>=3.0.5",
     "pytest>=8.2,<8.3.4",
@@ -108,6 +110,7 @@ gpu = ["cupy"]
 cu12 = ["cupy-cuda12x"]
 cu11 = ["cupy-cuda11x"]
 # https://github.com/dask/dask/issues/11290
+lazy = ["xarray>=2024.06.0", "aiohttp", "requests", "zarr<3.0.0a0", "anndata[dask]"]
 dask = ["dask[array]>=2022.09.2,!=2024.8.*,!=2024.9.*"]
 
 [tool.hatch.version]

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,27 @@
 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:  # pragma: no cover
+    msg = f"Cannot convert {type(anno)} to {attr} DataFrame"
+    raise ValueError(msg)
+
+
+@_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, cast
@@ -43,7 +43,6 @@
 from .sparse_dataset import BaseCompressedSparseDataset, sparse_dataset
 from .storage import coerce_array
 from .views import (
-    DataFrameView,
     DictView,
     _resolve_idxs,
     as_view,
@@ -55,7 +54,7 @@
     from typing import Any, Literal
 
     from ..compat import Index1D
-    from ..typing import ArrayDataStructureType
+    from ..typing import XDataType
     from .aligned_mapping import AxisArraysView, LayersView, PairwiseArraysView
     from .index import Index
 
@@ -291,8 +290,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
@@ -540,7 +539,7 @@ def shape(self) -> tuple[int, int]:
         return self.n_obs, self.n_vars
 
     @property
-    def X(self) -> ArrayDataStructureType | None:
+    def X(self) -> XDataType | None:
         """Data matrix of shape :attr:`n_obs` × :attr:`n_vars`."""
         if self.isbacked:
             if not self.file.is_open:
@@ -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

@@ -44,10 +44,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 DataArray
+
+    # 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
@@ -110,8 +113,12 @@
                 )
                 raise KeyError(msg)
             return positions  # np.ndarray[int]
+    elif isinstance(indexer, DataArray):
+        if isinstance(indexer.data, DaskArray):
+            return indexer.data.compute()
+        return indexer.data
     msg = f"Unknown indexer {indexer!r} of type {type(indexer)}"
-    raise IndexError(msg)
+    raise IndexError()
 
 
 def _fix_slice_bounds(s: slice, length: int) -> slice: