Backport PR #4446 on branch 1.12.x (fix: avoid loading whole dask arrays in axis_nnz and check_nonnegative_integers) (#4478)

Co-authored-by: Yangqi Su <yangqi.su@som.umaryland.edu>

Author: meeseeksmachine

docs/release-notes/4446.fix.md

@@ -0,0 +1 @@
+Fix {func}`~scanpy.pp.highly_variable_genes` (`flavor='seurat_v3'`), {func}`~scanpy.pp.calculate_qc_metrics`, and {func}`~scanpy.tl.rank_genes_groups` loading entire row-chunked {class}`dask.array.Array`s into memory when checking for count data or counting nonzero cells per gene {smaller}`Y Su`

src/scanpy/_utils/__init__.py

@@ -748,14 +748,32 @@ def _(x: CSBase, /, axis: Literal[0, 1]) -> np.ndarray:
 
 @axis_nnz.register(DaskArray)
 def _(x: DaskArray, /, axis: Literal[0, 1]) -> DaskArray:
-    return x.map_blocks(
-        partial(axis_nnz, axis=axis),
+    import dask.array as da
+
+    # Count per block, then sum the counts across blocks.
+    # (`map_blocks(..., drop_axis=axis)` would concatenate all blocks along `axis`
+    # before counting, i.e. load all of a row-chunked array into memory for `axis=0`.)
+    return da.reduction(
+        x,
+        chunk=_axis_nnz_block,
+        # keep the in-memory result’s dtype (int32 for scipy sparse), `np.sum` upcasts
+        aggregate=lambda counts, axis, keepdims: counts.sum(
+            axis=axis, keepdims=keepdims, dtype=counts.dtype
+        ),
+        axis=axis,
         dtype=np.int64,
         meta=np.array([], dtype=np.int64),
-        drop_axis=axis,
     )
 
 
+def _axis_nnz_block(
+    x: _MemoryArray, /, *, axis: tuple[Literal[0, 1]], keepdims: bool
+) -> np.ndarray:
+    # `da.reduction` passes `axis` as a tuple and expects the reduced axis to be kept
+    (ax,) = axis
+    return np.expand_dims(axis_nnz(x, axis=ax), ax)
+
+
 @singledispatch
 def check_nonnegative_integers(x: _SupportedArray, /) -> bool | DaskArray:
     """Check values of X to ensure it is count data."""
@@ -779,7 +797,18 @@ def _check_nonnegative_integers_in_mem(x: _MemoryArray, /) -> bool:
 
 @check_nonnegative_integers.register(DaskArray)
 def _check_nonnegative_integers_dask(x: DaskArray, /) -> DaskArray:
-    return x.map_blocks(check_nonnegative_integers, dtype=bool, drop_axis=(0, 1))
+    import dask.array as da
+
+    # One result per block, combined with `all` (see `axis_nnz` for why not `drop_axis`)
+    return da.reduction(
+        x,
+        chunk=lambda block, axis, keepdims: np.array([
+            [check_nonnegative_integers(block)]
+        ]),
+        aggregate=np.all,
+        dtype=bool,
+        meta=np.array([], dtype=bool),
+    )
 
 
 def dematrix[SA: _SupportedArray](x: SA | np.matrix) -> SA:

tests/test_utils.py

@@ -13,10 +13,12 @@
 from anndata.tests.helpers import asarray
 from scipy import sparse
 
+import scanpy._utils
 from scanpy._compat import CSBase, DaskArray
 from scanpy._utils import (
     _numba_thread_limit,
     axis_mul_or_truediv,
+    axis_nnz,
     check_nonnegative_integers,
     descend_classes_and_funcs,
 )
@@ -26,6 +28,7 @@
     random_k_tuples,
     random_str,
 )
+from testing.scanpy._pytest.marks import needs
 from testing.scanpy._pytest.params import (
     ARRAY_TYPES,
     ARRAY_TYPES_DASK,
@@ -182,6 +185,45 @@ def test_check_nonnegative_integers(array_type, array_value, expected):
     assert received is expected
 
 
+@needs.dask
+@pytest.mark.parametrize(
+    ("func", "kwargs"),
+    [
+        pytest.param(check_nonnegative_integers, {}, id="check_nonnegative_integers"),
+        pytest.param(axis_nnz, dict(axis=0), id="axis_nnz-0"),
+        pytest.param(axis_nnz, dict(axis=1), id="axis_nnz-1"),
+    ],
+)
+@pytest.mark.parametrize(
+    "block_type",
+    [np.asarray, sparse.csr_matrix],  # noqa: TID251
+    ids=["dense_dask", "csr_dask"],
+)
+def test_dask_reductions_blockwise(
+    monkeypatch: pytest.MonkeyPatch, func: Callable, kwargs: dict, block_type
+):
+    """Reductions over dask arrays must see one block at a time, not all blocks concatenated."""
+    import dask.array as da
+
+    x_np = np.random.default_rng(0).poisson(size=(100, 20)).astype(np.float32)
+    x = da.from_array(block_type(x_np), chunks=(10, 20), asarray=False)
+
+    # the dask implementation calls the module-level function on each block
+    seen_shapes = []
+
+    def spy(block, **kw):
+        seen_shapes.append(block.shape)
+        return func(block, **kw)
+
+    monkeypatch.setattr(scanpy._utils, func.__name__, spy)
+    lazy = func(x, **kwargs)
+    seen_shapes.clear()  # drop dask’s metadata probe, made while building the graph
+    received = lazy.compute()
+
+    assert seen_shapes == [(10, 20)] * 10
+    np.testing.assert_array_equal(received, func(x_np, **kwargs))
+
+
 @pytest.mark.parametrize("seed", [0, 1, 1256712675])
 @pytest.mark.parametrize("pass_seed", [True, False], ids=["pass_seed", "set_seed"])
 @pytest.mark.parametrize("func", ["choice"])