(feat): read_elem_as_dask method

src/anndata/_io/specs/__init__.py

@@ -1,21 +1,25 @@
 from __future__ import annotations
 
-from . import methods
+from . import lazy_methods, methods
 from .registry import (
+    _LAZY_REGISTRY,  # noqa: F401
     _REGISTRY,  # noqa: F401
     IOSpec,
     Reader,
     Writer,
     get_spec,
     read_elem,
+    read_elem_as_dask,
     write_elem,
 )
 
 __all__ = [
     "methods",
+    "lazy_methods",
     "write_elem",
     "get_spec",
     "read_elem",
+    "read_elem_as_dask",
     "Reader",
     "Writer",
     "IOSpec",

src/anndata/_io/specs/lazy_methods.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from pathlib import Path
+
+import h5py
+import numpy as np
+from scipy import sparse
+
+import anndata as ad
+from anndata.compat import H5Array, H5Group, ZarrArray, ZarrGroup
+
+from .registry import _LAZY_REGISTRY, IOSpec
+
+# TODO: settings
+stride = 100
+h5_chunks = 1000
+
+
+def make_dask_array(is_csc, shape, make_dask_chunk, dtype):
+    import dask.array as da
+
+    chunks = [None, None]
+    major_index = int(is_csc)
+    minor_index = (is_csc + 1) % 2
+    chunks[minor_index] = (shape[minor_index],)
+    chunks[major_index] = (stride,) * (shape[major_index] // stride) + (
+        shape[major_index] % stride,
+    )
+    memory_format = [sparse.csr_matrix, sparse.csc_matrix][major_index]
+    da_mtx = da.map_blocks(
+        make_dask_chunk,
+        dtype=dtype,
+        chunks=chunks,
+        meta=memory_format((0, 0), dtype=np.float32),
+    )
+    return da_mtx
+
+
+def make_index(is_csc, stride, shape, block_id):
+    index = (
+        slice(
+            block_id[is_csc] * stride,
+            min((block_id[is_csc] * stride) + stride, shape[0]),
+        ),
+    )
+    if is_csc:
+        return (slice(None, None, None),) + index
+    return index
+
+
+@contextmanager
+def maybe_open_h5(filename_or_elem: str | ZarrGroup, elem_name: str):
+    if isinstance(filename_or_elem, str):
+        file = h5py.File(filename_or_elem, "r")
+        try:
+            yield file[elem_name]
+        finally:
+            file.close()
+    else:
+        try:
+            yield filename_or_elem
+        finally:
+            pass
+
+
+@_LAZY_REGISTRY.register_read(H5Group, IOSpec("csc_matrix", "0.1.0"))
+@_LAZY_REGISTRY.register_read(H5Group, IOSpec("csr_matrix", "0.1.0"))
+@_LAZY_REGISTRY.register_read(ZarrGroup, IOSpec("csc_matrix", "0.1.0"))
+@_LAZY_REGISTRY.register_read(ZarrGroup, IOSpec("csr_matrix", "0.1.0"))
+def read_sparse_as_dask_h5(elem, _reader):
+    filename_or_elem = elem.file.filename if isinstance(elem, H5Group) else elem
+    elem_name = elem.name if isinstance(elem, H5Group) else Path(elem.path).name
+    shape = elem.attrs["shape"]
+    dtype = elem["data"].dtype
+    is_csc = elem.attrs["encoding-type"] == "csc_matrix"
+
+    def make_dask_chunk(block_id=None):
+        # We need to open the file in each task since `dask` cannot share h5py objects when using `dask.distributed`
+        # https://github.com/scverse/anndata/issues/1105
+        with maybe_open_h5(filename_or_elem, elem_name) as f:
+            mtx = ad.experimental.sparse_dataset(f)
+            index = make_index(is_csc, stride, shape, block_id)
+            chunk = mtx[index]
+        return chunk
+
+    return make_dask_array(is_csc, shape, make_dask_chunk, dtype)
+
+
+@_LAZY_REGISTRY.register_read(H5Array, IOSpec("array", "0.2.0"))
+def read_h5_array(elem, _reader):
+    import dask.array as da
+
+    if not hasattr(elem, "chunks") or elem.chunks is None:
+        return da.from_array(elem, chunks=(h5_chunks,) * len(elem.shape))
+    return da.from_array(elem)
+
+
+@_LAZY_REGISTRY.register_read(ZarrArray, IOSpec("array", "0.2.0"))
+def read_zarr_array(elem, _reader):
+    import dask.array as da
+
+    return da.from_zarr(elem)

src/anndata/_io/specs/registry.py

@@ -145,9 +145,7 @@ def get_reader(
         if (src_type, spec, modifiers) in self.read:
             return self.read[(src_type, spec, modifiers)]
         else:
-            raise IORegistryError._from_read_parts(
-                "read", _REGISTRY.read, src_type, spec
-            )
+            raise IORegistryError._from_read_parts("read", self.read, src_type, spec)
 
     def has_reader(
         self, src_type: type, spec: IOSpec, modifiers: frozenset[str] = frozenset()
@@ -176,7 +174,7 @@ def get_partial_reader(
             return self.read_partial[(src_type, spec, modifiers)]
         else:
             raise IORegistryError._from_read_parts(
-                "read_partial", _REGISTRY.read_partial, src_type, spec
+                "read_partial", self.read_partial, src_type, spec
             )
 
     def get_spec(self, elem: Any) -> IOSpec:
@@ -188,6 +186,7 @@ def get_spec(self, elem: Any) -> IOSpec:
 
 
 _REGISTRY = IORegistry()
+_LAZY_REGISTRY = IORegistry()
 
 
 @singledispatch
@@ -332,6 +331,21 @@ def read_elem(elem: StorageType) -> Any:
     return Reader(_REGISTRY).read_elem(elem)
 
 
+def read_elem_as_dask(elem: StorageType) -> Any:
+    """
+    Read an element from a store lazily.
+
+    Assumes that the element is encoded using the anndata encoding. This function will
+    determine the encoded type using the encoding metadata stored in elem's attributes.
+
+    Params
+    ------
+    elem
+        The stored element.
+    """
+    return Reader(_LAZY_REGISTRY).read_elem(elem)
+
+
 def write_elem(
     store: GroupStorageType,
     k: str,

tests/test_io_elementwise.py

@@ -15,7 +15,14 @@
 from scipy import sparse
 
 import anndata as ad
-from anndata._io.specs import _REGISTRY, IOSpec, get_spec, read_elem, write_elem
+from anndata._io.specs import (
+    _REGISTRY,
+    IOSpec,
+    get_spec,
+    read_elem,
+    read_elem_as_dask,
+    write_elem,
+)
 from anndata._io.specs.registry import IORegistryError
 from anndata.compat import H5Group, ZarrGroup, _read_attr
 from anndata.tests.helpers import (
@@ -47,6 +54,39 @@ def store(request, tmp_path) -> H5Group | ZarrGroup:
             file.close()
 
 
+sparse_formats = ["csr", "csc"]
+SIZE = 1000
+
+
+@pytest.fixture(scope="function", params=sparse_formats)
+def sparse_format(request):
+    return request.param
+
+
+def create_dense_store(store):
+    X = np.random.randn(SIZE, SIZE)
+
+    write_elem(store, "X", X)
+    return store
+
+
+def create_sparse_store(sparse_format, store):
+    import dask.array as da
+
+    X = sparse.random(
+        SIZE,
+        SIZE,
+        format=sparse_format,
+        density=0.01,
+        random_state=np.random.default_rng(),
+    )
+    X_dask = da.from_array(X, chunks=(100, 100))
+
+    write_elem(store, "X", X)
+    write_elem(store, "X_dask", X_dask)
+    return store
+
+
 @pytest.mark.parametrize(
     "value,encoding_type",
     [
@@ -126,30 +166,38 @@ def test_io_spec_cupy(store, value, encoding_type):
     assert get_spec(store[key]) == _REGISTRY.get_spec(value)
 
 
-@pytest.mark.parametrize("sparse_format", ["csr", "csc"])
-def test_dask_write_sparse(store, sparse_format):
-    import dask.array as da
+def test_dask_write_sparse(sparse_format, store):
+    x_sparse_store = create_sparse_store(sparse_format, store)
+    X_from_disk = read_elem(x_sparse_store["X"])
+    X_dask_from_disk = read_elem(x_sparse_store["X_dask"])
 
-    X = sparse.random(
-        1000,
-        1000,
-        format=sparse_format,
-        density=0.01,
-        random_state=np.random.default_rng(),
-    )
-    X_dask = da.from_array(X, chunks=(100, 100))
+    assert_equal(X_from_disk, X_dask_from_disk)
+    assert_equal(dict(x_sparse_store["X"].attrs), dict(x_sparse_store["X_dask"].attrs))
 
-    write_elem(store, "X", X)
-    write_elem(store, "X_dask", X_dask)
+    assert x_sparse_store["X_dask/indptr"].dtype == np.int64
+    assert x_sparse_store["X_dask/indices"].dtype == np.int64
 
-    X_from_disk = read_elem(store["X"])
-    X_dask_from_disk = read_elem(store["X_dask"])
+
+@pytest.mark.parametrize("arr_type", ["dense", *sparse_formats])
+def test_read_lazy_2d_dask(arr_type, store):
+    if arr_type == "dense":
+        arr_store = create_dense_store(store)
+    else:
+        arr_store = create_sparse_store(arr_type, store)
+    X_dask_from_disk = read_elem_as_dask(arr_store["X"])
+    X_from_disk = read_elem(arr_store["X"])
 
     assert_equal(X_from_disk, X_dask_from_disk)
-    assert_equal(dict(store["X"].attrs), dict(store["X_dask"].attrs))
+    random_int_indices = np.random.randint(0, SIZE, (SIZE // 10,))
+    random_bool_mask = np.random.randn(SIZE) > 0
+    index_slice = slice(0, SIZE // 10)
+    for index in [random_int_indices, index_slice, random_bool_mask]:
+        assert_equal(X_from_disk[index, :], X_dask_from_disk[index, :])
+        assert_equal(X_from_disk[:, index], X_dask_from_disk[:, index])
 
-    assert store["X_dask/indptr"].dtype == np.int64
-    assert store["X_dask/indices"].dtype == np.int64
+    if arr_type in {"csr", "csc"}:
+        assert arr_store["X_dask/indptr"].dtype == np.int64
+        assert arr_store["X_dask/indices"].dtype == np.int64
 
 
 @pytest.mark.parametrize("sparse_format", ["csr", "csc"])
@@ -195,7 +243,7 @@ def test_write_anndata_to_root(store):
     ["attribute", "value"],
     [
         ("encoding-type", "floob"),
-        ("encoding-version", "10000.0"),
+        ("encoding-version", "SIZE0.0"),
     ],
 )
 def test_read_iospec_not_found(store, attribute, value):