(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')
 

.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

@@ -21,7 +21,7 @@
 - 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 {attr}`~anndata.settings.shall_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 :func:`~anndata.experimental.read_elem_lazy` 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.shall_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_backed_experimental
 ```

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,6 +50,8 @@
     | pd.Index,
     index: pd.Index,
 ) -> slice | int | np.ndarray:  # ndarray of int or bool
+    from ..experimental.backed._compat import xr
+
     if not isinstance(index, pd.RangeIndex):
         msg = "Don’t call _normalize_index with non-categorical/string names"
         assert index.dtype != float, msg
@@ -107,6 +109,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)}")
 

src/anndata/_core/merge.py

@@ -757,6 +757,7 @@ def np_bool_to_pd_bool_array(df: pd.DataFrame):
     return df
 
 
+# TODO: concat for xarray
 def concat_arrays(arrays, reindexers, axis=0, index=None, fill_value=None):
     arrays = list(arrays)
     if fill_value is None:

src/anndata/_core/storage.py

@@ -25,6 +25,16 @@
     allow_df: bool = False,
     allow_array_like: bool = False,
 ):
+    try:
+        # Needs to be done here to prevent circular imports, and StorageType is immutable
+        from anndata.experimental.backed._xarray import Dataset2D
+    except ImportError:
+
+        class Dataset2D:
+            @staticmethod
+            def __repr__():
+                return "mock anndata.experimental.backed._xarray."
+
     """Coerce arrays stored in layers/X, and aligned arrays ({obs,var}{m,p})."""
     from ..typing import ArrayDataStructureType
 
@@ -33,7 +43,7 @@
         return value
     # If value is one of the allowed types, return it
     array_data_structure_types = get_args(ArrayDataStructureType)
-    if isinstance(value, array_data_structure_types):
+    if isinstance(value, (*array_data_structure_types, Dataset2D)):
         if isinstance(value, np.matrix):
             msg = f"{name} should not be a np.matrix, use np.ndarray instead."
             warnings.warn(msg, ImplicitModificationWarning)

src/anndata/_core/views.py

@@ -13,6 +13,7 @@
 
 from anndata._warnings import ImplicitModificationWarning
 
+from .._settings import settings
 from ..compat import (
     AwkArray,
     CupyArray,
@@ -305,6 +306,11 @@ def as_view_dask_array(array, view_args):
 
 @as_view.register(pd.DataFrame)
 def as_view_df(df, view_args):
+    if settings.shall_remove_unused_categories:
+        for col in df.columns:
+            if isinstance(df[col].dtype, pd.CategoricalDtype):
+                with pd.option_context("mode.chained_assignment", None):
+                    df[col] = df[col].cat.remove_unused_categories()
     return DataFrameView(df, view_args=view_args)
 
 

src/anndata/_io/specs/__init__.py

@@ -9,7 +9,7 @@
     Writer,
     get_spec,
     read_elem,
-    read_elem_as_dask,
+    read_elem_lazy,
     write_elem,
 )
 
@@ -19,7 +19,7 @@
     "write_elem",
     "get_spec",
     "read_elem",
-    "read_elem_as_dask",
+    "read_elem_lazy",
     "Reader",
     "Writer",
     "IOSpec",