(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

@@ -121,7 +121,8 @@ Low level methods for reading and writing elements of an `AnnData` object to a s
 
    experimental.read_elem
    experimental.write_elem
-   experimental.read_elem_as_dask
+   experimental.read_elem_lazy
+   experimental.read_backed
 ```
 
 Utilities for customizing the IO process:
@@ -148,6 +149,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"),
 ]

docs/release-notes/0.11.0.md

@@ -6,9 +6,9 @@
 * Allow `axis` parameter of e.g. :func:`anndata.concat` to accept `'obs'` and `'var` {pr}`1244` {user}`flying-sheep`
 * Add `settings` object with methods for altering internally-used options, like checking for uniqueness on `obs`' index {pr}`1270` {user}`ilan-gold`
 * `scipy.sparse.csr_array` and `scipy.sparse.csc_array` are now supported when constructing `AnnData` objects {pr}`1028` {user}`ilan-gold` {user}`isaac-virshup`
+* Add {func}`~anndata.experimental.read_elem_lazy` to handle backed dataframes, sparse arrays, and dense arrays, as well as a {func}`~anndata.experimental.read_backed` to handle reading in as much of the {class}`~anndata.AnnData` as possible backed {pr}`1469` {pr}`1247` {user}`ilan-gold`
 * Add {attr}`~anndata.settings.should_remove_unused_categories` option to {attr}`anndata.settings` to override current behavior {pr}`1340` {user}`ilan-gold`
 * Add {attr}`~anndata.settings.should_check_uniqueness` option to {attr}`anndata.settings` to override current behavior {pr}`1507` {user}`ilan-gold`
-* Add :func:`~anndata.experimental.read_elem_as_dask` function to handle i/o with sparse and dense arrays {pr}`1469` {user}`ilan-gold`
 * Add functionality to write from GPU {class}`dask.array.Array` to disk {pr}`1550` {user}`ilan-gold`
 * 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 {pr}`1474` {user}` falexwolf`
 

pyproject.toml

@@ -98,12 +98,15 @@ test = [
     "awkward>=2.3",
     "pyarrow",
     "pytest_memray",
-    "pytest-mock"
+    "pytest-mock",
 ]
+test-full = ["anndata[test]", "anndata[xarray]"]
 gpu = [
     "cupy",
     "numpy<2.0.0",
 ]
+xarray = ["xarray>=2024.06.0"]
+
 
 [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
@@ -36,7 +36,6 @@
 from .sparse_dataset import BaseCompressedSparseDataset, sparse_dataset
 from .storage import coerce_array
 from .views import (
-    DataFrameView,
     DictView,
     _resolve_idxs,
     as_view,
@@ -80,6 +79,32 @@
         )
 
 
+@singledispatch
+def _remove_unused_categories(
+    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):
+            continue
+        all_categories = df_full[k].cat.categories
+        with pd.option_context("mode.chained_assignment", None):
+            df_sub[k] = df_sub[k].cat.remove_unused_categories()
+        # also correct the colors...
+        color_key = f"{k}_colors"
+        if color_key not in uns:
+            continue
+        color_vec = uns[color_key]
+        if np.array(color_vec).ndim == 0:
+            # Make 0D arrays into 1D ones
+            uns[color_key] = np.array(color_vec)[(None,)]
+        elif len(color_vec) != len(all_categories):
+            # Reset colors
+            del uns[color_key]
+        else:
+            idx = np.where(np.isin(all_categories, df_sub[k].cat.categories))[0]
+            uns[color_key] = np.array(color_vec)[(idx,)]
+
+
 class AnnData(metaclass=utils.DeprecationMixinMeta):
     """\
     An annotated data matrix.
@@ -289,11 +314,11 @@
         # fix categories
         uns = copy(adata_ref._uns)
         if settings.should_remove_unused_categories:
-            self._remove_unused_categories(adata_ref.obs, obs_sub, uns)
-            self._remove_unused_categories(adata_ref.var, var_sub, uns)
+            _remove_unused_categories(adata_ref.obs, obs_sub, uns)
+            _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
@@ -1013,30 +1038,6 @@
         oidx, vidx = self._normalize_indices(index)
         return AnnData(self, oidx=oidx, vidx=vidx, asview=True)
 
-    def _remove_unused_categories(
-        self, 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):
-                continue
-            all_categories = df_full[k].cat.categories
-            with pd.option_context("mode.chained_assignment", None):
-                df_sub[k] = df_sub[k].cat.remove_unused_categories()
-            # also correct the colors...
-            color_key = f"{k}_colors"
-            if color_key not in uns:
-                continue
-            color_vec = uns[color_key]
-            if np.array(color_vec).ndim == 0:
-                # Make 0D arrays into 1D ones
-                uns[color_key] = np.array(color_vec)[(None,)]
-            elif len(color_vec) != len(all_categories):
-                # Reset colors
-                del uns[color_key]
-            else:
-                idx = np.where(np.isin(all_categories, df_sub[k].cat.categories))[0]
-                uns[color_key] = np.array(color_vec)[(idx,)]
-
     def rename_categories(self, key: str, categories: Sequence[Any]):
         """\
         Rename categories of annotation `key` in :attr:`obs`, :attr:`var`,

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/sparse_dataset.py

@@ -27,9 +27,8 @@
 from scipy.sparse import _sparsetools
 
 from anndata._core.index import _fix_slice_bounds
-from anndata.compat import H5Group, ZarrArray, ZarrGroup
 
-from ..compat import SpArray, _read_attr
+from ..compat import H5Group, SpArray, ZarrArray, ZarrGroup, _read_attr
 
 try:
     # Not really important, just for IDEs to be more helpful

src/anndata/_core/storage.py

@@ -72,12 +72,23 @@
     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})."""
     # If value is a scalar and we allow that, return it
     if allow_array_like and np.isscalar(value):
         return value
     # If value is one of the allowed types, return it
-    if isinstance(value, StorageType.classes()):
+
+    if isinstance(value, StorageType.classes()) or isinstance(value, 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.should_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",