Allow xarray Datasets to be used for obs/var/obsm/varm

.github/workflows/test-cpu.yml

@@ -62,7 +62,7 @@ jobs:
         if: matrix.dependencies-version == 'minimum'
         run: |
           uv pip install --system --compile tomli packaging
-          deps=$(python3 ci/scripts/min-deps.py pyproject.toml --extra dev test)
+          deps=$(python3 ci/scripts/min-deps.py pyproject.toml --extra dev test-min)
           uv pip install --system --compile $deps "anndata @ ."
 
       - name: Install dependencies release candidates

ci/scripts/min-deps.py

@@ -34,7 +34,7 @@ def min_dep(req: Requirement) -> Requirement:
     -------
 
     >>> min_dep(Requirement("numpy>=1.0"))
-    <Requirement('numpy==1.0.*')>
+    <Requirement('numpy~=1.0.0')>
     >>> min_dep(Requirement("numpy<3.0"))
     <Requirement('numpy<3.0')>
     """
@@ -55,7 +55,7 @@ def min_dep(req: Requirement) -> Requirement:
         elif spec.operator == "==":
             min_version = Version(spec.version)
 
-    return Requirement(f"{req_name}=={min_version}.*")
+    return Requirement(f"{req_name}~={min_version}.0")
 
 
 def extract_min_deps(
@@ -64,6 +64,7 @@ def extract_min_deps(
     dependencies = deque(dependencies)  # We'll be mutating this
     project_name = pyproject["project"]["name"]
 
+    deps = {}
     while len(dependencies) > 0:
         req = dependencies.pop()
 
@@ -76,7 +77,11 @@ def extract_min_deps(
                 extra_deps = pyproject["project"]["optional-dependencies"][extra]
                 dependencies += map(Requirement, extra_deps)
         else:
-            yield min_dep(req)
+            if req.name in deps:
+                req.specifier &= deps[req.name].specifier
+                req.extras |= deps[req.name].extras
+            deps[req.name] = min_dep(req)
+    yield from deps.values()
 
 
 class Args(argparse.Namespace):

docs/api.md

@@ -180,7 +180,7 @@ Types used by the former:
    experimental.StorageType
    experimental.backed._lazy_arrays.MaskedArray
    experimental.backed._lazy_arrays.CategoricalArray
-   experimental.backed._xarray.Dataset2D
+   _core.xarray.Dataset2D
 ```
 
 (extensions-api)=

docs/conf.py

@@ -142,6 +142,7 @@ def setup(app: Sphinx):
     "anndata.compat.DaskArray": "dask.array.Array",
     "anndata.compat.CupyArray": "cupy.ndarray",
     "anndata.compat.CupySparseMatrix": "cupyx.scipy.sparse.spmatrix",
+    "anndata.compat.XDataArray": "xarray.DataArray",
     "awkward.highlevel.Array": "ak.Array",
     "numpy.int64": ("py:attr", "numpy.int64"),
     "pandas.DataFrame.iloc": ("py:attr", "pandas.DataFrame.iloc"),

docs/release-notes/0.12.0rc1.md

@@ -10,7 +10,7 @@
 
 #### Bug fixes
 
-- Disallow writing of {class}`~anndata.experimental.backed._xarray.Dataset2D` objects {user}`ilan-gold` ({pr}`1887`)
+- Disallow writing of {class}`~anndata._core.xarray.Dataset2D` objects {user}`ilan-gold` ({pr}`1887`)
 - Upgrade old deprecation warning to a `FutureWarning` on `BaseCompressedSparseDataset.__setitem__`, showing our intent to remove the feature in the next release.  {user}`ilan-gold` ({pr}`1928`)
 - Don't use {func}`asyncio.run` internally for any operations {user}`ilan-gold` ({pr}`1933`)
 - Disallow forward slashes in keys for writing {user}`ilan-gold` ({pr}`1940`)

docs/release-notes/1966.feature.md

@@ -0,0 +1 @@
+Allow xarray Datasets to be used for obs/var/obsm/varm. {user}`ilia-kats`

pyproject.toml

@@ -79,11 +79,9 @@ doc = [
     "sphinx_design>=0.5.0",
     # for unreleased changes
     "anndata[dev-doc,dask]",
-    "awkward>=2.3",
 ]
 dev-doc = [ "towncrier>=24.8.0" ] # release notes tool
-test-full = [ "anndata[test,lazy]" ]
-test = [
+test-min = [
     "loompy>=3.0.5",
     "pytest>=8.2,<8.3.4",
     "pytest-cov",
@@ -100,15 +98,19 @@ test = [
     "scanpy>=1.10",
     "httpx",                # For data downloading
     "dask[distributed]",
-    "awkward>=2.3",
+    "awkward>=2.3.2",
     "pyarrow",
     "anndata[dask]",
 ]
+test = [
+    "anndata[test-min,lazy]",
+    "pandas>=2.1.0",
+] # pandas 2.1.0 needs to be specified for xarray to work with min-deps script
 gpu = [ "cupy" ]
 cu12 = [ "cupy-cuda12x" ]
 cu11 = [ "cupy-cuda11x" ]
 # requests and aiohttp needed for zarr remote data
-lazy = [ "xarray>=2024.06.0", "aiohttp", "requests", "anndata[dask]" ]
+lazy = [ "xarray>=2025.04.0", "aiohttp", "requests", "anndata[dask]" ]
 # https://github.com/dask/dask/issues/11290
 # https://github.com/dask/dask/issues/11752
 dask = [ "dask[array]>=2023.5.1,!=2024.8.*,!=2024.9.*,<2025.2.0" ]

src/anndata/_core/aligned_df.py

@@ -9,6 +9,8 @@
 from pandas.api.types import is_string_dtype
 
 from .._warnings import ImplicitModificationWarning
+from ..compat import XDataset
+from .xarray import Dataset2D
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
@@ -108,15 +110,39 @@ def _mk_df_error(
     expected: int,
     actual: int,
 ):
+    what = "row" if attr == "obs" else "column"
     if source == "X":
-        what = "row" if attr == "obs" else "column"
         msg = (
             f"Observations annot. `{attr}` must have as many rows as `X` has {what}s "
             f"({expected}), but has {actual} rows."
         )
     else:
         msg = (
             f"`shape` is inconsistent with `{attr}` "
-            "({actual} {what}s instead of {expected})"
+            f"({actual} {what}s instead of {expected})"
         )
     return ValueError(msg)
+
+
+@_gen_dataframe.register(Dataset2D)
+def _gen_dataframe_xr(
+    anno: Dataset2D,
+    index_names: Iterable[str],
+    *,
+    source: Literal["X", "shape"],
+    attr: Literal["obs", "var"],
+    length: int | None = None,
+):
+    return anno
+
+
+@_gen_dataframe.register(XDataset)
+def _gen_dataframe_xdataset(
+    anno: XDataset,
+    index_names: Iterable[str],
+    *,
+    source: Literal["X", "shape"],
+    attr: Literal["obs", "var"],
+    length: int | None = None,
+):
+    return Dataset2D(anno)

src/anndata/_core/aligned_mapping.py

@@ -11,7 +11,7 @@
 import pandas as pd
 
 from .._warnings import ExperimentalFeatureWarning, ImplicitModificationWarning
-from ..compat import AwkArray, CSArray, CSMatrix, CupyArray
+from ..compat import AwkArray, CSArray, CSMatrix, CupyArray, XDataset
 from ..utils import (
     axis_len,
     convert_to_dict,
@@ -23,6 +23,7 @@
 from .index import _subset
 from .storage import coerce_array
 from .views import as_view, view_update
+from .xarray import Dataset2D
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Iterable, Iterator, Mapping
@@ -75,8 +76,10 @@ def _validate_value(self, val: Value, key: str) -> Value:
                 ExperimentalFeatureWarning,
                 # stacklevel=3,
             )
-        if isinstance(val, np.ndarray | CupyArray) and len(val.shape) == 1:
+        elif isinstance(val, np.ndarray | CupyArray) and len(val.shape) == 1:
             val = val.reshape((val.shape[0], 1))
+        elif isinstance(val, XDataset):
+            val = Dataset2D(data_vars=val.data_vars, coords=val.coords, attrs=val.attrs)
         for i, axis in enumerate(self.axes):
             if self.parent.shape[axis] == axis_len(val, i):
                 continue
@@ -275,6 +278,9 @@ def _validate_value(self, val: Value, key: str) -> Value:
                 else:
                     msg = "Index.equals and pd.testing.assert_index_equal disagree"
                     raise AssertionError(msg)
+            val.index.name = (
+                self.dim_names.name
+            )  # this is consistent with AnnData.obsm.setter and AnnData.varm.setter
         return super()._validate_value(val, key)
 
     @property

src/anndata/_core/anndata.py

@@ -47,6 +47,7 @@
     _resolve_idxs,
     as_view,
 )
+from .xarray import Dataset2D
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
@@ -55,7 +56,7 @@
 
     from zarr.storage import StoreLike
 
-    from ..compat import Index1D
+    from ..compat import Index1D, XDataset
     from ..typing import XDataType
     from .aligned_mapping import AxisArraysView, LayersView, PairwiseArraysView
     from .index import Index
@@ -746,10 +747,14 @@ def n_vars(self) -> int:
         """Number of variables/features."""
         return len(self.var_names)
 
-    def _set_dim_df(self, value: pd.DataFrame, attr: Literal["obs", "var"]):
-        if not isinstance(value, pd.DataFrame):
-            msg = f"Can only assign pd.DataFrame to {attr}."
-            raise ValueError(msg)
+    def _set_dim_df(self, value: pd.DataFrame | XDataset, attr: Literal["obs", "var"]):
+        value = _gen_dataframe(
+            value,
+            [f"{attr}_names", f"{'row' if attr == 'obs' else 'col'}_names"],
+            source="shape",
+            attr=attr,
+            length=self.n_obs if attr == "obs" else self.n_vars,
+        )
         raise_value_error_if_multiindex_columns(value, attr)
         value_idx = self._prep_dim_index(value.index, attr)
         if self.is_view:
@@ -804,12 +809,12 @@ def _set_dim_index(self, value: pd.Index, attr: str):
                 v.index = value
 
     @property
-    def obs(self) -> pd.DataFrame:
+    def obs(self) -> pd.DataFrame | Dataset2D:
         """One-dimensional annotation of observations (`pd.DataFrame`)."""
         return self._obs
 
     @obs.setter
-    def obs(self, value: pd.DataFrame):
+    def obs(self, value: pd.DataFrame | XDataset):
         self._set_dim_df(value, "obs")
 
     @obs.deleter
@@ -827,12 +832,12 @@ def obs_names(self, names: Sequence[str]):
         self._set_dim_index(names, "obs")
 
     @property
-    def var(self) -> pd.DataFrame:
+    def var(self) -> pd.DataFrame | Dataset2D:
         """One-dimensional annotation of variables/ features (`pd.DataFrame`)."""
         return self._var
 
     @var.setter
-    def var(self, value: pd.DataFrame):
+    def var(self, value: pd.DataFrame | XDataset):
         self._set_dim_df(value, "var")
 
     @var.deleter
@@ -2079,6 +2084,14 @@ def _get_and_delete_multicol_field(self, a, key_multicol):
         return values
 
 
+@AnnData._remove_unused_categories.register(Dataset2D)
+@staticmethod
+def _remove_unused_categories_xr(
+    df_full: Dataset2D, df_sub: Dataset2D, uns: dict[str, Any]
+):
+    pass  # this is handled automatically by the categorical arrays themselves i.e., they dedup upon access.
+
+
 def _check_2d_shape(X):
     """\
     Check shape of array or sparse matrix.

src/anndata/_core/file_backing.py

@@ -10,6 +10,7 @@
 
 from ..compat import AwkArray, DaskArray, ZarrArray, ZarrGroup
 from .sparse_dataset import BaseCompressedSparseDataset
+from .xarray import Dataset2D
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
@@ -162,6 +163,11 @@ def _(x: AwkArray, *, copy: bool = False):
         return x
 
 
+@to_memory.register(Dataset2D)
+def _(x: Dataset2D, *, copy: bool = False):
+    return x.to_memory(copy=copy)
+
+
 @singledispatch
 def filename(x):
     msg = f"Not implemented for {type(x)}"

src/anndata/_core/index.py

@@ -10,7 +10,8 @@
 import pandas as pd
 from scipy.sparse import issparse
 
-from ..compat import AwkArray, CSArray, CSMatrix, DaskArray
+from ..compat import AwkArray, CSArray, CSMatrix, DaskArray, XDataArray
+from .xarray import Dataset2D
 
 if TYPE_CHECKING:
     from ..compat import Index, Index1D
@@ -44,8 +45,6 @@ def _normalize_index(  # noqa: PLR0911, PLR0912
     | pd.Index,
     index: pd.Index,
 ) -> slice | int | np.ndarray:  # ndarray of int or bool
-    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) and index.dtype in (np.float64, np.int64):
         msg = f"Don’t call _normalize_index with non-categorical/string names and non-range index {index}"
@@ -112,7 +111,7 @@ def name_idx(i):
                 )
                 raise KeyError(msg)
             return positions  # np.ndarray[int]
-    elif isinstance(indexer, DataArray):
+    elif isinstance(indexer, XDataArray):
         if isinstance(indexer.data, DaskArray):
             return indexer.data.compute()
         return indexer.data
@@ -210,6 +209,15 @@ def _subset_awkarray(a: AwkArray, subset_idx: Index):
     return a[subset_idx]
 
 
+@_subset.register(Dataset2D)
+def _(a: Dataset2D, subset_idx: Index):
+    key = a.index_dim
+    # xarray seems to have some code looking for a second entry in tuples
+    if isinstance(subset_idx, tuple) and len(subset_idx) == 1:
+        subset_idx = subset_idx[0]
+    return a.isel(**{key: subset_idx})
+
+
 # Registration for SparseDataset occurs in sparse_dataset.py
 @_subset.register(h5py.Dataset)
 def _subset_dataset(d, subset_idx):