Add sum_where to duck_array_ops

Author: tomwhite

xarray/core/duck_array_ops.py

@@ -281,6 +281,16 @@ def count(data, axis=None):
     return np.sum(np.logical_not(isnull(data)), axis=axis)
 
 
+def sum_where(data, axis=None, dtype=None, where=None):
+    xp = get_array_namespace(data)
+    if where is not None:
+        a = where_method(xp.zeros_like(data), where, data)
+    else:
+        a = data
+    result = xp.sum(a, axis=axis, dtype=dtype)
+    return result
+
+
 def where(condition, x, y):
     """Three argument where() with better dtype promotion rules."""
     xp = get_array_namespace(condition)

xarray/core/nanops.py

@@ -5,14 +5,7 @@
 import numpy as np
 
 from . import dtypes, nputils, utils
-from .duck_array_ops import (
-    count,
-    fillna,
-    get_array_namespace,
-    isnull,
-    where,
-    where_method,
-)
+from .duck_array_ops import count, fillna, isnull, sum_where, where, where_method
 
 
 def _replace_nan(a, val):
@@ -24,7 +17,6 @@ def _replace_nan(a, val):
     return where_method(val, mask, a), mask
 
 
-
 def _maybe_null_out(result, axis, mask, min_count=1):
     """
     xarray version of pandas.core.nanops._maybe_null_out
@@ -100,13 +92,8 @@ def nanargmax(a, axis=None):
 
 
 def nansum(a, axis=None, dtype=None, out=None, min_count=None):
-    if hasattr(a, "__array_namespace__"):
-        a, mask = _replace_nan(a, 0)
-        xp = get_array_namespace(a)
-        result = xp.sum(a, axis=axis, dtype=dtype)
-    else:
-        mask = isnull(a)
-        result = np.nansum(a, axis=axis, dtype=dtype)      
+    mask = isnull(a)
+    result = sum_where(a, axis=axis, dtype=dtype, where=mask)
     if min_count is not None:
         return _maybe_null_out(result, axis, mask, min_count)
     else:

xarray/tests/test_array_api.py

@@ -17,8 +17,16 @@
 
 @pytest.fixture
 def arrays() -> tuple[xr.DataArray, xr.DataArray]:
-    np_arr = xr.DataArray(np.ones((2, 3)), dims=("x", "y"), coords={"x": [10, 20]})
-    xp_arr = xr.DataArray(xp.ones((2, 3)), dims=("x", "y"), coords={"x": [10, 20]})
+    np_arr = xr.DataArray(
+        np.array([[1.0, 2.0, 3.0], [4.0, 5.0, np.nan]]),
+        dims=("x", "y"),
+        coords={"x": [10, 20]},
+    )
+    xp_arr = xr.DataArray(
+        xp.asarray([[1.0, 2.0, 3.0], [4.0, 5.0, np.nan]]),
+        dims=("x", "y"),
+        coords={"x": [10, 20]},
+    )
     assert isinstance(xp_arr.data, Array)
     return np_arr, xp_arr