feat: Support sparrays in dask (#2087)

Author: flying-sheep

docs/release-notes/2087.feature.md

@@ -0,0 +1 @@
+Add support for Dask {class}`dask.array.Array`s containing {class}`~scipy.sparse.sparray`s {user}`flying-sheep`

pyproject.toml

@@ -101,7 +101,7 @@ test-min = [
     "pyarrow<21",        # https://github.com/scikit-hep/awkward/issues/3579
     "anndata[dask]",
 ]
-test = [ "anndata[test-min,lazy]" ]
+test = [ "anndata[test-min,lazy]", "fast-array-utils>=1.2.3" ]
 gpu = [ "cupy" ]
 cu12 = [ "cupy-cuda12x" ]
 cu11 = [ "cupy-cuda11x" ]

src/anndata/_core/merge.py

@@ -146,7 +146,7 @@ def equal_dask_array(a, b) -> bool:
         return False
     if isinstance(b, DaskArray) and tokenize(a) == tokenize(b):
         return True
-    if isinstance(a._meta, CSMatrix):
+    if isinstance(a._meta, CSMatrix | CSArray):
         # TODO: Maybe also do this in the other case?
         return da.map_blocks(equal, a, b, drop_axis=(0, 1)).all()
     else:

src/anndata/_io/specs/methods.py

@@ -876,6 +876,8 @@ def chunk_slice(start: int, stop: int) -> tuple[slice | None, slice | None]:
     for scipy_sparse_type, spec in [
         (sparse.csr_matrix, IOSpec("csr_matrix", "0.1.0")),
         (sparse.csc_matrix, IOSpec("csc_matrix", "0.1.0")),
+        (sparse.csr_array, IOSpec("csr_matrix", "0.1.0")),
+        (sparse.csc_array, IOSpec("csc_matrix", "0.1.0")),
     ]:
         _REGISTRY.register_write(group_type, (array_type, scipy_sparse_type), spec)(
             write_dask_sparse

src/anndata/compat/__init__.py

@@ -11,7 +11,7 @@
 import h5py
 import numpy as np
 import pandas as pd
-import scipy
+import scipy.sparse
 from numpy.typing import NDArray
 from packaging.version import Version
 from zarr import Array as ZarrArray  # noqa: F401

src/anndata/tests/helpers.py

@@ -6,6 +6,7 @@
 from collections import Counter, defaultdict
 from collections.abc import Mapping
 from functools import partial, singledispatch, wraps
+from importlib.metadata import version
 from importlib.util import find_spec
 from string import ascii_letters
 from typing import TYPE_CHECKING
@@ -14,6 +15,7 @@
 import numpy as np
 import pandas as pd
 import pytest
+from packaging.version import Version
 from pandas.api.types import is_numeric_dtype
 from scipy import sparse
 
@@ -61,6 +63,14 @@
         *(pd.UInt8Dtype, pd.UInt16Dtype, pd.UInt32Dtype, pd.UInt64Dtype),
     )
 
+try:
+    import fast_array_utils as _
+except ImportError:
+    # dask natively supports sparray since https://github.com/dask/dask/pull/11750
+    DASK_CAN_SPARRAY = Version(version("dask")) >= Version("2025.3.0")
+else:  # fast-array-utils monkeypatches dask to support sparrays
+    DASK_CAN_SPARRAY = True
+
 
 DEFAULT_KEY_TYPES = (
     sparse.csr_matrix,
@@ -628,8 +638,9 @@ def assert_equal_arrayview(
 
 @assert_equal.register(BaseCompressedSparseDataset)
 @assert_equal.register(sparse.spmatrix)
+@assert_equal.register(CSArray)
 def assert_equal_sparse(
-    a: BaseCompressedSparseDataset | sparse.spmatrix,
+    a: BaseCompressedSparseDataset | sparse.spmatrix | CSArray,
     b: object,
     *,
     exact: bool = False,
@@ -639,13 +650,6 @@ def assert_equal_sparse(
     assert_equal(b, a, exact=exact, elem_name=elem_name)
 
 
-@assert_equal.register(CSArray)
-def assert_equal_sparse_array(
-    a: CSArray, b: object, *, exact: bool = False, elem_name: str | None = None
-):
-    return assert_equal_sparse(a, b, exact=exact, elem_name=elem_name)
-
-
 @assert_equal.register(CupySparseMatrix)
 def assert_equal_cupy_sparse(
     a: CupySparseMatrix, b: object, *, exact: bool = False, elem_name: str | None = None
@@ -878,29 +882,37 @@ def _(a):
 
 
 @singledispatch
-def as_sparse_dask_array(a) -> DaskArray:
-    import dask.array as da
-
-    return da.from_array(sparse.csr_matrix(a), chunks=_half_chunk_size(a.shape))
+def _as_sparse_dask(
+    a: NDArray | CSArray | CSMatrix | DaskArray, *, typ: type[CSArray | CSMatrix]
+) -> DaskArray:
+    """Convert a to a sparse dask array, preserving sparse format and container (`cs{rc}_{array,matrix}`)."""
+    raise NotImplementedError
 
 
-@as_sparse_dask_array.register(CSMatrix)
-def _(a):
+@_as_sparse_dask.register(CSArray | CSMatrix | np.ndarray)
+def _(a: CSArray | CSMatrix | NDArray, *, typ: type[CSArray | CSMatrix]) -> DaskArray:
     import dask.array as da
 
-    return da.from_array(a, _half_chunk_size(a.shape))
+    return da.from_array(_as_sparse_dask_inner(a, typ=typ), _half_chunk_size(a.shape))
 
 
-@as_sparse_dask_array.register(CSArray)
-def _(a):
-    import dask.array as da
+@_as_sparse_dask.register(DaskArray)
+def _(a: DaskArray, *, typ: type[CSArray | CSMatrix]) -> DaskArray:
+    return a.map_blocks(_as_sparse_dask_inner, typ=typ, dtype=a.dtype, meta=typ((2, 2)))
 
-    return da.from_array(sparse.csr_matrix(a), _half_chunk_size(a.shape))
 
+def _as_sparse_dask_inner(
+    a: NDArray | CSArray | CSMatrix, *, typ: type[CSArray | CSMatrix]
+) -> CSArray | CSMatrix:
+    """Convert into a a sparse container that dask supports (or complain)."""
+    if issubclass(typ, CSArray) and not DASK_CAN_SPARRAY:  # convert sparray to spmatrix
+        msg = "Dask <2025.3 without fast-array-utils doesn’t support sparse arrays"
+        raise TypeError(msg)
+    return typ(a)
 
-@as_sparse_dask_array.register(DaskArray)
-def _(a):
-    return a.map_blocks(sparse.csr_matrix)
+
+as_sparse_dask_array = partial(_as_sparse_dask, typ=sparse.csr_array)
+as_sparse_dask_matrix = partial(_as_sparse_dask, typ=sparse.csr_matrix)
 
 
 @singledispatch
@@ -949,11 +961,8 @@ def _(a):
 # We should try and fix this upstream in dask/ cupy
 @singledispatch
 def as_cupy_sparse_dask_array(a, format="csr"):
-    memory_class = format_to_memory_class[format]
-    cpu_da = as_sparse_dask_array(a)
-    return cpu_da.rechunk((cpu_da.chunks[0], -1)).map_blocks(
-        memory_class, dtype=a.dtype, meta=memory_class(cpu_da._meta)
-    )
+    da = _as_sparse_dask(a, typ=format_to_memory_class[format])
+    return da.rechunk((da.chunks[0], -1))
 
 
 @as_cupy_sparse_dask_array.register(CupyArray)
@@ -1003,7 +1012,7 @@ def as_cupy(val, typ=None):
     if issubclass(typ, CupyArray):
         import cupy as cp
 
-        if isinstance(val, CSMatrix):
+        if isinstance(val, CSMatrix | CSArray):
             val = val.toarray()
         return cp.array(val)
     elif issubclass(typ, CupyCSRMatrix):
@@ -1059,7 +1068,14 @@ def shares_memory_sparse(x, y):
 
 DASK_MATRIX_PARAMS = [
     pytest.param(as_dense_dask_array, id="dense_dask_array"),
-    pytest.param(as_sparse_dask_array, id="sparse_dask_array"),
+    pytest.param(as_sparse_dask_matrix, id="sparse_dask_matrix"),
+    pytest.param(
+        as_sparse_dask_array,
+        marks=pytest.mark.skipif(
+            not DASK_CAN_SPARRAY, reason="Dask does not support sparrays"
+        ),
+        id="sparse_dask_array",
+    ),
 ]
 
 CUPY_MATRIX_PARAMS = [

tests/test_dask.py

@@ -4,31 +4,37 @@
 
 from __future__ import annotations
 
+from pathlib import Path
 from typing import TYPE_CHECKING
 
 import numpy as np
 import pandas as pd
 import pytest
-from scipy import sparse
 
 import anndata as ad
 from anndata._core.anndata import AnnData
-from anndata.compat import CupyArray, DaskArray
+from anndata.compat import CSArray, CSMatrix, CupyArray, DaskArray
 from anndata.experimental.merge import as_group
 from anndata.tests.helpers import (
+    BASE_MATRIX_PARAMS,
+    DASK_CAN_SPARRAY,
+    DASK_MATRIX_PARAMS,
     GEN_ADATA_DASK_ARGS,
-    as_cupy_sparse_dask_array,
     as_dense_cupy_dask_array,
     as_dense_dask_array,
     as_sparse_dask_array,
+    as_sparse_dask_matrix,
     assert_equal,
     gen_adata,
 )
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
     from pathlib import Path
     from typing import Literal
 
+    from numpy.typing import NDArray
+
 
 pytest.importorskip("dask.array")
 
@@ -276,13 +282,18 @@ def test_assign_X(adata):
 @pytest.mark.parametrize(
     ("array_func", "mem_type"),
     [
-        pytest.param(as_dense_dask_array, np.ndarray, id="dense_dask_array"),
-        pytest.param(as_sparse_dask_array, sparse.csr_matrix, id="sparse_dask_array"),
+        pytest.param(as_dense_dask_array, np.ndarray, id="dense"),
+        pytest.param(as_sparse_dask_matrix, CSMatrix, id="sparse_matrix"),
+        pytest.param(
+            as_sparse_dask_array,
+            CSArray,
+            marks=pytest.mark.skipif(
+                not DASK_CAN_SPARRAY, reason="Dask does not support sparrays"
+            ),
+            id="sparse_array",
+        ),
         pytest.param(
-            as_dense_cupy_dask_array,
-            CupyArray,
-            id="cupy_dense_dask_array",
-            marks=pytest.mark.gpu,
+            as_dense_cupy_dask_array, CupyArray, id="cupy_dense", marks=pytest.mark.gpu
         ),
     ],
 )
@@ -311,29 +322,17 @@ def test_dask_to_memory_unbacked(array_func, mem_type):
     assert isinstance(orig.uns["da"]["da"], DaskArray)
 
 
-@pytest.mark.parametrize(
-    "array_func",
-    [
-        pytest.param(as_dense_dask_array, id="dense_dask_array"),
-        pytest.param(as_sparse_dask_array, id="sparse_dask_array"),
-        pytest.param(
-            as_dense_cupy_dask_array,
-            id="cupy_dense_dask_array",
-            marks=pytest.mark.gpu,
-        ),
-        pytest.param(
-            as_cupy_sparse_dask_array,
-            id="cupy_sparse_dask_array",
-            marks=pytest.mark.gpu,
-        ),
-    ],
-)
-def test_dask_to_disk_view(array_func, diskfmt, tmp_path):
+@pytest.mark.parametrize("to_dask", [*BASE_MATRIX_PARAMS, *DASK_MATRIX_PARAMS])
+def test_dask_to_disk_view(
+    to_dask: Callable[[NDArray], DaskArray],
+    diskfmt: Literal["h5ad", "zarr"],
+    tmp_path: Path,
+) -> None:
     random_state = np.random.default_rng()
-    orig = ad.AnnData(
-        # need to change type for cupy
-        array_func(random_state.binomial(100, 0.005, (20, 15)).astype("float32"))
-    )
+    arr = random_state.binomial(100, 0.005, (20, 15)).astype("float32")
+
+    # TODO: need to change type for cupy
+    orig = ad.AnnData(to_dask(arr))
     orig = orig[orig.shape[0] // 2]
     path = tmp_path / f"test.{diskfmt}"
     getattr(orig, f"write_{diskfmt}")(path)

tests/test_io_elementwise.py

@@ -252,7 +252,7 @@ def test_io_spec_compressed_scalars(store: G, value: np.ndarray, encoding_type:
 @pytest.mark.parametrize("as_dask", [False, True])
 def test_io_spec_cupy(store, value, encoding_type, as_dask):
     if as_dask:
-        if isinstance(value, CSMatrix):
+        if isinstance(value, CSMatrix | CSArray):
             value = as_cupy_sparse_dask_array(value, format=encoding_type[:3])
         else:
             value = as_dense_cupy_dask_array(value)