Backport PR #2034: fix: no cache indptr zarr dask (#2045)

Co-authored-by: Ilan Gold <[email protected]>

Author: meeseeksmachine

benchmarks/benchmarks/sparse_dataset.py

@@ -4,11 +4,13 @@
 
 import numpy as np
 import zarr
+from dask.array.core import Array as DaskArray
 from scipy import sparse
 
 from anndata import AnnData
 from anndata._core.sparse_dataset import sparse_dataset
 from anndata._io.specs import write_elem
+from anndata.experimental import read_elem_lazy
 
 
 def make_alternating_mask(n):
@@ -37,27 +39,36 @@ class SparseCSRContiguousSlice:
             # (10_000, 500)
         ],
         _slices.keys(),
+        [True, False],
     )
-    param_names = ("shape", "slice")
+    param_names = ("shape", "slice", "use_dask")
 
-    def setup(self, shape: tuple[int, int], slice: str):
+    def setup(self, shape: tuple[int, int], slice: str, use_dask: bool):  # noqa: FBT001
         X = sparse.random(
             *shape, density=0.01, format="csr", random_state=np.random.default_rng(42)
         )
         self.slice = self._slices[slice]
         g = zarr.group()
         write_elem(g, "X", X)
-        self.x = sparse_dataset(g["X"])
+        self.x = read_elem_lazy(g["X"]) if use_dask else sparse_dataset(g["X"])
         self.adata = AnnData(self.x)
 
     def time_getitem(self, *_):
-        self.x[self.slice]
+        res = self.x[self.slice]
+        if isinstance(res, DaskArray):
+            res.compute()
 
     def peakmem_getitem(self, *_):
-        self.x[self.slice]
+        res = self.x[self.slice]
+        if isinstance(res, DaskArray):
+            res.compute()
 
     def time_getitem_adata(self, *_):
-        self.adata[self.slice]
+        res = self.adata[self.slice]
+        if isinstance(res, DaskArray):
+            res.compute()
 
     def peakmem_getitem_adata(self, *_):
-        self.adata[self.slice]
+        res = self.adata[self.slice]
+        if isinstance(res, DaskArray):
+            res.compute()

src/anndata/_core/sparse_dataset.py

@@ -165,7 +165,11 @@ def _offsets(
     def _get_contiguous_compressed_slice(
         self, s: slice
     ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-        new_indptr = self.indptr[s.start : s.stop + 1].copy()
+        new_indptr = self.indptr[s.start : s.stop + 1]
+        # If indptr is cached, we need to make a copy of the subset
+        # so as not to alter the underlying cached data.
+        if isinstance(self.indptr, np.ndarray):
+            new_indptr = new_indptr.copy()
 
         start = new_indptr[0]
         stop = new_indptr[-1]

src/anndata/_io/specs/lazy_methods.py

@@ -132,7 +132,7 @@ def read_sparse_as_dask(
     path_or_sparse_dataset = (
         Path(filename(elem))
         if isinstance(elem, H5Group)
-        else ad.io.sparse_dataset(elem)
+        else ad.io.sparse_dataset(elem, should_cache_indptr=False)
     )
     elem_name = get_elem_name(elem)
     shape: tuple[int, int] = tuple(elem.attrs["shape"])