pydata / pydata/xarray

`sum(..., skipna=True)` should dispatch to `sparse.nansum` for sparse arrays

Open
#10,922 9 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

topic-arrays topic-performance
Dominant language
Python
Stars
4.2k
Forks
1.4k
Avg merge
2d 15h
Merged PRs (30d)
14

Description

What is your issue?

xarray currently uses its own nanops.nansum when calling DataArray.sum(..., skipna=None), which relies on sum_where. This implementation route is very inefficient for sparse arrays, especially (and ironically) when operating on a sparse array with fill_value=np.nan, see https://github.com/pydata/sparse/issues/908. Why doesn't xarray try to dispatch to a possible nansum implementation in the array's namespace?

sparse offers its own nansum. Internally, it also seems to use where, but it's much faster than the xarray nansum. I applied the following patch to duck_array_ops.py, reducing the time for sums on a sparse array significantly:

--- duck_array_ops.py	2025-11-14 12:21:49
+++ duck_array_ops.py	2025-11-14 12:23:20
@@ -519,6 +519,15 @@
 
             nanname = "nan" + name
             func = getattr(nanops, nanname)
+
+            if "min_count" not in kwargs or kwargs["min_count"] is None:
+                try:
+                    kwargs.pop("min_count", None)
+                    xp = get_array_namespace(values)
+                    func = getattr(xp, name)
+                except AttributeError:
+                    pass
+
         else:
             if name in ["sum", "prod"]:
                 kwargs.pop("min_count", None)

Dispatching to sparse.nansum produces a factor 20+ speedup:

# Without patch
$ !python -m timeit -s "import sparse; import xarray as xr; import numpy as np; arr = xr.DataArray(sparse.random((100, 100, 100), density=0.1, fill_value=np.nan), dims=['x', 'y', 'z'])" "arr.sum(dim=['y', 'z'])"
1 loop, best of 5: 36.2 msec per loop

# With patch
$ !python -m timeit -s "import sparse; import xarray as xr; import numpy as np; arr = xr.DataArray(sparse.random((100, 100, 100), density=0.1, fill_value=np.nan), dims=['x', 'y', 'z'])" "arr.sum(dim=['y', 'z'])"
200 loops, best of 5: 1.37 msec per loop

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

The dispatch logic is in duck_array_ops.py; start by reading the handling of nansum, min_count, and get_array_namespace. Use the provided sparse-array timeit benchmark to compare behavior and performance. Done means sums on sparse arrays can use sparse.nansum without breaking the existing min_count cases.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
data, performance
Issue type
Feature
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.