(feat): experimental read_backed method for zarr + hdf5 via read_dispatched

.gitignore

@@ -1,6 +1,7 @@
 # Temp files
 .DS_Store
 *~
+venv/
 
 # Compiled files
 __pycache__/

anndata/_core/aligned_mapping.py

@@ -43,7 +43,7 @@ class AlignedMapping(cabc.MutableMapping, ABC):
     """The actual class (which has it’s own data) for this aligned mapping."""
 
     def __repr__(self):
-        return f"{type(self).__name__} with keys: {', '.join(self.keys())}"
+        return f"{type(self).__name__} with keys: {', '.join([k + f'[{str(self[k].dtype)}]' for k in self.keys()])}"
 
     def _ipython_key_completions_(self) -> List[str]:
         return list(self.keys())
@@ -119,7 +119,9 @@ def copy(self):
 
     def _view(self, parent: "anndata.AnnData", subset_idx: I):
         """Returns a subset copy-on-write view of the object."""
-        return self._view_class(self, parent, subset_idx)
+        if parent.is_view or subset_idx is not None:  # and or or?
+            return self._view_class(self, parent, subset_idx)
+        return self
 
     @deprecated("dict(obj)")
     def as_dict(self) -> dict:

anndata/_core/anndata.py

@@ -46,7 +46,7 @@
     as_view,
     _resolve_idxs,
 )
-from .sparse_dataset import SparseDataset
+from .sparse_dataset import sparse_dataset
 from .. import utils
 from ..utils import convert_to_dict, ensure_df_homogeneous, dim_len
 from ..logging import anndata_logger as logger
@@ -58,6 +58,7 @@
     CupySparseMatrix,
     _move_adj_mtx,
 )
+from .sparse_dataset import BaseCompressedSparseDataset
 
 
 class StorageType(Enum):
@@ -67,6 +68,7 @@ class StorageType(Enum):
     ZarrArray = ZarrArray
     ZappyArray = ZappyArray
     DaskArray = DaskArray
+    BaseCompressedSparseDataset = BaseCompressedSparseDataset
     CupyArray = CupyArray
     CupySparseMatrix = CupySparseMatrix
 
@@ -609,11 +611,13 @@ def X(self) -> Optional[Union[np.ndarray, sparse.spmatrix, ArrayView]]:
                 self.file.open()
             X = self.file["X"]
             if isinstance(X, h5py.Group):
-                X = SparseDataset(X)
+                X = sparse_dataset(X)
             # This is so that we can index into a backed dense dataset with
             # indices that aren’t strictly increasing
             if self.is_view:
                 X = _subset(X, (self._oidx, self._vidx))
+                if isinstance(X, BaseCompressedSparseDataset):
+                    X = X.to_memory()
         elif self.is_view and self._adata_ref.X is None:
             X = None
         elif self.is_view:
@@ -674,7 +678,7 @@ def X(self, value: Optional[Union[np.ndarray, sparse.spmatrix]]):
                 if self.is_view:
                     X = self.file["X"]
                     if isinstance(X, h5py.Group):
-                        X = SparseDataset(X)
+                        X = sparse_dataset(X)
                     X[oidx, vidx] = value
                 else:
                     self._set_backed("X", value)

anndata/_core/file_backing.py

@@ -7,8 +7,8 @@
 import h5py
 
 from . import anndata
-from .sparse_dataset import SparseDataset
-from ..compat import ZarrArray, DaskArray, AwkArray
+from .sparse_dataset import BaseCompressedSparseDataset
+from ..compat import ZarrArray, ZarrGroup, DaskArray, AwkArray
 
 
 class AnnDataFileManager:
@@ -39,11 +39,15 @@ def __contains__(self, x) -> bool:
     def __iter__(self) -> Iterator[str]:
         return iter(self._file)
 
-    def __getitem__(self, key: str) -> Union[h5py.Group, h5py.Dataset, SparseDataset]:
+    def __getitem__(
+        self, key: str
+    ) -> Union[h5py.Group, h5py.Dataset, BaseCompressedSparseDataset]:
         return self._file[key]
 
     def __setitem__(
-        self, key: str, value: Union[h5py.Group, h5py.Dataset, SparseDataset]
+        self,
+        key: str,
+        value: Union[h5py.Group, h5py.Dataset, BaseCompressedSparseDataset],
     ):
         self._file[key] = value
 
@@ -110,8 +114,8 @@ def _(x, copy=False):
     return x[...]
 
 
-@to_memory.register(SparseDataset)
-def _(x: SparseDataset, copy=False):
+@to_memory.register(BaseCompressedSparseDataset)
+def _(x: BaseCompressedSparseDataset, copy=True):
     return x.to_memory()
 
 
@@ -133,3 +137,20 @@ def _(x, copy=False):
         return _copy(x)
     else:
         return x
+
+
+@singledispatch
+def filename(x):
+    raise NotImplementedError(f"Not implemented for {type(x)}")
+
+
+@filename.register(h5py.Group)
+@filename.register(h5py.Dataset)
+def _(x):
+    return x.file.filename
+
+
+@filename.register(ZarrArray)
+@filename.register(ZarrGroup)
+def _(x):
+    return x.store.path

anndata/_core/index.py

@@ -7,7 +7,7 @@
 import numpy as np
 import pandas as pd
 from scipy.sparse import spmatrix, issparse
-from ..compat import AwkArray, DaskArray, Index, Index1D
+from ..compat import AwkArray, DaskArray, Index, Index1D, ZarrArray
 
 
 def _normalize_indices(
@@ -116,16 +116,25 @@ def unpack_index(index: Index) -> Tuple[Index1D, Index1D]:
 
 
 @singledispatch
-def _subset(a: Union[np.ndarray, pd.DataFrame], subset_idx: Index):
+def _subset(a: np.ndarray, subset_idx: Index):
     # Select as combination of indexes, not coordinates
     # Correcting for indexing behaviour of np.ndarray
     if all(isinstance(x, cabc.Iterable) for x in subset_idx):
         subset_idx = np.ix_(*subset_idx)
     return a[subset_idx]
 
 
+@_subset.register(ZarrArray)
+def _subset_zarr(a: ZarrArray, subset_idx: Index):
+    if all(isinstance(x, cabc.Iterable) for x in subset_idx):
+        subset_idx = np.ix_(*subset_idx)
+    return a.oindex[subset_idx]
+
+
 @_subset.register(DaskArray)
 def _subset_dask(a: DaskArray, subset_idx: Index):
+    if isinstance(subset_idx, slice):
+        return a[subset_idx]
     if all(isinstance(x, cabc.Iterable) for x in subset_idx):
         subset_idx = np.ix_(*subset_idx)
         return a.vindex[subset_idx]
@@ -147,6 +156,8 @@ def _subset_df(df: pd.DataFrame, subset_idx: Index):
 
 @_subset.register(AwkArray)
 def _subset_awkarray(a: AwkArray, subset_idx: Index):
+    if isinstance(subset_idx, slice):
+        return a[subset_idx]
     if all(isinstance(x, cabc.Iterable) for x in subset_idx):
         subset_idx = np.ix_(*subset_idx)
     return a[subset_idx]

anndata/_core/raw.py

@@ -9,7 +9,7 @@
 from . import anndata
 from .index import _normalize_index, _subset, unpack_index, get_vector
 from .aligned_mapping import AxisArrays
-from .sparse_dataset import SparseDataset
+from .sparse_dataset import BaseCompressedSparseDataset, sparse_dataset
 
 from ..compat import CupyArray, CupySparseMatrix
 
@@ -49,7 +49,7 @@ def _get_X(self, layer=None):
         return self.X
 
     @property
-    def X(self) -> Union[SparseDataset, np.ndarray, sparse.spmatrix]:
+    def X(self) -> Union[BaseCompressedSparseDataset, np.ndarray, sparse.spmatrix]:
         # TODO: Handle unsorted array of integer indices for h5py.Datasets
         if not self._adata.isbacked:
             return self._X
@@ -66,7 +66,7 @@ def X(self) -> Union[SparseDataset, np.ndarray, sparse.spmatrix]:
                 f"{self._adata.file.filename}."
             )
         if isinstance(X, h5py.Group):
-            X = SparseDataset(X)
+            X = sparse_dataset(X)
         # Check if we need to subset
         if self._adata.is_view:
             # TODO: As noted above, implement views of raw
@@ -187,6 +187,7 @@ class _RawViewHack:
     def __init__(self, raw: Raw, vidx: Union[slice, np.ndarray]):
         self.parent_raw = raw
         self.vidx = vidx
+        self.is_view = True
 
     @property
     def shape(self) -> Tuple[int, int]: