(feat): async full reading of sparse

Author: ilan-gold

src/anndata/_core/sparse_dataset.py

@@ -12,6 +12,7 @@
 # - think about supporting the COO format
 from __future__ import annotations
 
+import asyncio
 import warnings
 from abc import ABC
 from collections.abc import Iterable
@@ -661,6 +662,31 @@ def to_memory(self) -> CSMatrix | CSArray:
         mtx.indptr = self._indptr
         return mtx
 
+    async def to_memory_async(self) -> CSMatrix | CSArray:
+        format_class = get_memory_class(
+            self.format, use_sparray_in_io=settings.use_sparse_array_on_read
+        )
+        mtx = format_class(self.shape, dtype=self.dtype)
+        mtx.indptr = self._indptr
+        if isinstance(self._data, ZarrArray):
+            await asyncio.gather(
+                *(
+                    self.set_memory_async_from_zarr(mtx, attr)
+                    for attr in ["data", "indices"]
+                )
+            )
+        else:
+            mtx.data = self._data[...]
+            mtx.indices = self._indices[...]
+        return mtx
+
+    async def set_memory_async_from_zarr(
+        self, mtx: CSMatrix | CSArray, attr: Literal["indptr", "data", "indices"]
+    ) -> None:
+        setattr(
+            mtx, attr, await getattr(self, f"_{attr}")._async_array.getitem(())
+        )  # TODO: better way to asyncify
+
 
 class _CSRDataset(BaseCompressedSparseDataset, abc.CSRDataset):
     """Internal concrete version of :class:`anndata.abc.CSRDataset`."""

src/anndata/_io/specs/methods.py

@@ -831,7 +831,7 @@ def chunk_slice(start: int, stop: int) -> tuple[slice | None, slice | None]:
 @_REGISTRY.register_read(ZarrGroup, IOSpec("csc_matrix", "0.1.0"))
 @_REGISTRY.register_read(ZarrGroup, IOSpec("csr_matrix", "0.1.0"))
 async def read_sparse(elem: GroupStorageType, *, _reader: Reader) -> CSMatrix | CSArray:
-    return sparse_dataset(elem).to_memory()
+    return await sparse_dataset(elem).to_memory_async()
 
 
 #################