(feat): read_elem_as_dask method

docs/_templates/autosummary/class.rst

@@ -13,7 +13,7 @@
    .. autosummary::
       :toctree: .
    {% for item in attributes %}
-      ~{{ fullname }}.{{ item }}
+      ~{{ name }}.{{ item }}
    {%- endfor %}
    {% endif %}
    {% endblock %}
@@ -26,7 +26,7 @@
       :toctree: .
    {% for item in methods %}
       {%- if item != '__init__' %}
-      ~{{ fullname }}.{{ item }}
+      ~{{ name }}.{{ item }}
       {%- endif -%}
    {%- endfor %}
    {% endif %}

docs/api.md

@@ -121,6 +121,7 @@ Low level methods for reading and writing elements of an `AnnData` object to a s
 
    experimental.read_elem
    experimental.write_elem
+   experimental.read_elem_as_dask
 ```
 
 Utilities for customizing the IO process:

docs/conf.py

@@ -27,7 +27,7 @@
 # default settings
 templates_path = ["_templates"]
 html_static_path = ["_static"]
-source_suffix = [".rst", ".md"]
+source_suffix = {".rst": "restructuredtext", ".md": "markdown"}
 master_doc = "index"
 default_role = "literal"
 exclude_patterns = [

docs/release-notes/0.11.0.md

@@ -8,6 +8,7 @@
 * Add `should_remove_unused_categories` option to `anndata.settings` to override current behavior.  Default is `True` (i.e., previous behavior).  Please refer to the [documentation](https://anndata.readthedocs.io/en/latest/generated/anndata.settings.html) for usage.  {pr}`1340` {user}`ilan-gold`
 * `scipy.sparse.csr_array` and `scipy.sparse.csc_array` are now supported when constructing `AnnData` objects {pr}`1028` {user}`ilan-gold` {user}`isaac-virshup`
 * Add `should_check_uniqueness` option to `anndata.settings` to override current behavior.  Default is `True` (i.e., previous behavior).  Please refer to the [documentation](https://anndata.readthedocs.io/en/latest/generated/anndata.settings.html) for usage.  {pr}`1507` {user}`ilan-gold`
+* Add :func:`~anndata.experimental.read_elem_as_dask` function to handle i/o with sparse and dense arrays {pr}`1469` {user}`ilan-gold`
 * Add functionality to write from GPU {class}`dask.array.Array` to disk {pr}`1550` {user}`ilan-gold`
 
 #### Bugfix

pyproject.toml

@@ -66,7 +66,7 @@ dev = [
     "pytest-xdist",
 ]
 doc = [
-    "sphinx>=4.4",
+    "sphinx>=7.4.6",
     "sphinx-book-theme>=1.1.0",
     "sphinx-autodoc-typehints>=2.2.0",
     "sphinx-issues",

src/anndata/_core/file_backing.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Mapping
 from functools import singledispatch
-from pathlib import Path
+from pathlib import Path, PurePosixPath
 from typing import TYPE_CHECKING
 
 import h5py
@@ -161,3 +161,18 @@ def _(x):
 @filename.register(ZarrGroup)
 def _(x):
     return x.store.path
+
+
+@singledispatch
+def get_elem_name(x):
+    raise NotImplementedError(f"Not implemented for {type(x)}")
+
+
+@get_elem_name.register(h5py.Group)
+def _(x):
+    return x.name
+
+
+@get_elem_name.register(ZarrGroup)
+def _(x):
+    return PurePosixPath(x.path).name

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,169 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import h5py
+import numpy as np
+from scipy import sparse
+
+import anndata as ad
+from anndata._core.file_backing import filename, get_elem_name
+from anndata.compat import H5Array, H5Group, ZarrArray, ZarrGroup
+
+from .registry import _LAZY_REGISTRY, IOSpec
+
+if TYPE_CHECKING:
+    from typing import Literal, Union
+
+    from anndata.compat import DaskArray
+
+    from .registry import DaskReader
+
+
+@contextmanager
+def maybe_open_h5(path_or_group: Path | ZarrGroup, elem_name: str):
+    if not isinstance(path_or_group, Path):
+        yield path_or_group
+        return
+    file = h5py.File(path_or_group, "r")
+    try:
+        yield file[elem_name]
+    finally:
+        file.close()
+
+
+_DEFAULT_STRIDE = 1000
+
+
+def compute_chunk_layout_for_axis_shape(
+    chunk_axis_shape: int, full_axis_shape: int
+) -> tuple[int, ...]:
+    n_strides, rest = np.divmod(full_axis_shape, chunk_axis_shape)
+    chunk = (chunk_axis_shape,) * n_strides
+    if rest > 0:
+        chunk += (rest,)
+    return chunk
+
+
+@_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(
+    elem: H5Group | ZarrGroup,
+    *,
+    _reader: DaskReader,
+    chunks: tuple[int, ...] | None = None,
+) -> DaskArray:
+    import dask.array as da
+
+    path_or_group = Path(filename(elem)) if isinstance(elem, H5Group) else elem
+    elem_name = get_elem_name(elem)
+    shape: tuple[int, int] = tuple(elem.attrs["shape"])
+    dtype = elem["data"].dtype
+    is_csc: bool = elem.attrs["encoding-type"] == "csc_matrix"
+
+    stride: int = _DEFAULT_STRIDE
+    if chunks is not None:
+        if len(chunks) != 2:
+            raise ValueError("`chunks` must be a tuple of two integers")
+        if chunks[int(not is_csc)] != shape[int(not is_csc)]:
+            raise ValueError(
+                "Only the major axis can be chunked. "
+                f"Try setting chunks to {((-1, _DEFAULT_STRIDE) if is_csc else (_DEFAULT_STRIDE, -1))}"
+            )
+        stride = chunks[int(is_csc)]
+
+    def make_dask_chunk(
+        block_info: Union[  # noqa: UP007
+            dict[
+                Literal[None],
+                dict[str, Union[tuple[int, ...], list[tuple[int, ...]]]],  # noqa: UP007
+            ],
+            None,
+        ] = 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
+        if block_info is None:
+            raise ValueError("Block info is required")
+        with maybe_open_h5(path_or_group, elem_name) as f:
+            mtx = ad.experimental.sparse_dataset(f)
+            array_location = block_info[None]["array-location"]
+            index = (
+                slice(array_location[0][0], array_location[0][1]),
+                slice(array_location[1][0], array_location[1][1]),
+            )
+            chunk = mtx[index]
+        return chunk
+
+    shape_minor, shape_major = shape if is_csc else shape[::-1]
+    chunks_major = compute_chunk_layout_for_axis_shape(stride, shape_major)
+    chunks_minor = (shape_minor,)
+    chunk_layout = (
+        (chunks_minor, chunks_major) if is_csc else (chunks_major, chunks_minor)
+    )
+    memory_format = sparse.csc_matrix if is_csc else sparse.csr_matrix
+    da_mtx = da.map_blocks(
+        make_dask_chunk,
+        dtype=dtype,
+        chunks=chunk_layout,
+        meta=memory_format((0, 0), dtype=dtype),
+    )
+    return da_mtx
+
+
+@_LAZY_REGISTRY.register_read(H5Array, IOSpec("array", "0.2.0"))
+def read_h5_array(
+    elem: H5Array, *, _reader: DaskReader, chunks: tuple[int, ...] | None = None
+) -> DaskArray:
+    import dask.array as da
+
+    path = Path(elem.file.filename)
+    elem_name = elem.name
+    shape = tuple(elem.shape)
+    dtype = elem.dtype
+    chunks: tuple[int, ...] = (
+        chunks if chunks is not None else (_DEFAULT_STRIDE,) * len(shape)
+    )
+
+    def make_dask_chunk(
+        block_info: Union[  # noqa: UP007
+            dict[
+                Literal[None],
+                dict[str, Union[tuple[int, ...], list[tuple[int, ...]]]],  # noqa: UP007
+            ],
+            None,
+        ] = None,
+    ):
+        if block_info is None:
+            raise ValueError("Block info is required")
+        with maybe_open_h5(path, elem_name) as f:
+            idx = ()
+            for i in range(len(shape)):
+                array_location = block_info[None]["array-location"][i]
+                idx += (slice(array_location[0], array_location[1]),)
+            return f[idx]
+
+    chunk_layout = tuple(
+        compute_chunk_layout_for_axis_shape(chunks[i], shape[i])
+        for i in range(len(shape))
+    )
+
+    return da.map_blocks(
+        make_dask_chunk,
+        dtype=dtype,
+        chunks=chunk_layout,
+    )
+
+
+@_LAZY_REGISTRY.register_read(ZarrArray, IOSpec("array", "0.2.0"))
+def read_zarr_array(
+    elem: ZarrArray, *, _reader: DaskReader, chunks: tuple[int, ...] | None = None
+) -> DaskArray:
+    chunks: tuple[int, ...] = chunks if chunks is not None else elem.chunks
+    import dask.array as da
+
+    return da.from_zarr(elem, chunks=chunks)