Fix behaviour of min_count in reducing functions

doc/whats-new.rst

@@ -113,6 +113,11 @@ Bug fixes
   By `Leif Denby <https://github.com/leifdenby>`_.
 - Fix time encoding bug associated with using cftime versions greater than
   1.4.0 with xarray (:issue:`4870`, :pull:`4871`). By `Spencer Clark <https://github.com/spencerkclark>`_.
+- Stop ``sum`` and ``prod`` computing lazy arrays when called with a ``min_count``
+  parameter (:issue:`4898`, :pull:`4911`). By `Blair Bonnett <https://github.com/bcbnz>`_.
+- Fix bug preventing the ``min_count`` parameter to ``sum`` and ``prod`` working correctly
+  when calculating over all axes of a float64 array (:issue:`4898`, :pull:`4911`).
+  By `Blair Bonnett <https://github.com/bcbnz>`_.
 
 Documentation
 ~~~~~~~~~~~~~

xarray/core/dtypes.py

@@ -78,7 +78,7 @@ def maybe_promote(dtype):
     return np.dtype(dtype), fill_value
 
 
-NAT_TYPES = (np.datetime64("NaT"), np.timedelta64("NaT"))
+NAT_TYPES = {np.datetime64("NaT").dtype, np.timedelta64("NaT").dtype}
 
 
 def get_fill_value(dtype):

xarray/core/nanops.py

@@ -28,18 +28,14 @@ def _maybe_null_out(result, axis, mask, min_count=1):
     """
     xarray version of pandas.core.nanops._maybe_null_out
     """
-
     if axis is not None and getattr(result, "ndim", False):
         null_mask = (np.take(mask.shape, axis).prod() - mask.sum(axis) - min_count) < 0
-        if null_mask.any():
-            dtype, fill_value = dtypes.maybe_promote(result.dtype)
-            result = result.astype(dtype)
-            result[null_mask] = fill_value
+        dtype, fill_value = dtypes.maybe_promote(result.dtype)
+        result = np.where(null_mask, fill_value, result.astype(dtype))
 
     elif getattr(result, "dtype", None) not in dtypes.NAT_TYPES:
         null_mask = mask.size - mask.sum()
-        if null_mask < min_count:
-            result = np.nan
+        result = np.where(null_mask < min_count, np.nan, result)
 
     return result
 

xarray/tests/test_dtypes.py

@@ -90,3 +90,9 @@ def test_maybe_promote(kind, expected):
     actual = dtypes.maybe_promote(np.dtype(kind))
     assert actual[0] == expected[0]
     assert str(actual[1]) == expected[1]
+
+
+def test_nat_types_membership():
+    assert np.datetime64("NaT").dtype in dtypes.NAT_TYPES
+    assert np.timedelta64("NaT").dtype in dtypes.NAT_TYPES
+    assert np.float64 not in dtypes.NAT_TYPES

xarray/tests/test_duck_array_ops.py

@@ -34,6 +34,7 @@
     assert_array_equal,
     has_dask,
     has_scipy,
+    raise_if_dask_computes,
     raises_regex,
     requires_cftime,
     requires_dask,
@@ -587,7 +588,10 @@ def test_min_count(dim_num, dtype, dask, func, aggdim, contains_nan, skipna):
     da = construct_dataarray(dim_num, dtype, contains_nan=contains_nan, dask=dask)
     min_count = 3
 
-    actual = getattr(da, func)(dim=aggdim, skipna=skipna, min_count=min_count)
+    # If using Dask, the function call should be lazy.
+    with raise_if_dask_computes():
+        actual = getattr(da, func)(dim=aggdim, skipna=skipna, min_count=min_count)
+
     expected = series_reduce(da, func, skipna=skipna, dim=aggdim, min_count=min_count)
     assert_allclose(actual, expected)
     assert_dask_array(actual, dask)
@@ -603,14 +607,62 @@ def test_min_count_nd(dtype, dask, func):
     min_count = 3
     dim_num = 3
     da = construct_dataarray(dim_num, dtype, contains_nan=True, dask=dask)
-    actual = getattr(da, func)(dim=["x", "y", "z"], skipna=True, min_count=min_count)
+
+    # If using Dask, the function call should be lazy.
+    with raise_if_dask_computes():
+        actual = getattr(da, func)(
+            dim=["x", "y", "z"], skipna=True, min_count=min_count
+        )
+
     # Supplying all dims is equivalent to supplying `...` or `None`
     expected = getattr(da, func)(dim=..., skipna=True, min_count=min_count)
 
     assert_allclose(actual, expected)
     assert_dask_array(actual, dask)
 
 
+@pytest.mark.parametrize("dask", [False, True])
+@pytest.mark.parametrize("func", ["sum", "prod"])
+@pytest.mark.parametrize("dim", [None, "a", "b"])
+def test_min_count_specific(dask, func, dim):
+    if dask and not has_dask:
+        pytest.skip("requires dask")
+
+    # Simple array with four non-NaN values.
+    da = DataArray(np.ones((6, 6), dtype=np.float64) * np.nan, dims=("a", "b"))
+    da[0][0] = 2
+    da[0][3] = 2
+    da[3][0] = 2
+    da[3][3] = 2
+    if dask:
+        da = da.chunk({"a": 3, "b": 3})
+
+    # Expected result if we set min_count to the number of non-NaNs in a
+    # row/column/the entire array.
+    if dim:
+        min_count = 2
+        expected = DataArray(
+            [4.0, np.nan, np.nan] * 2, dims=("a" if dim == "b" else "b",)
+        )
+    else:
+        min_count = 4
+        expected = DataArray(8.0 if func == "sum" else 16.0)
+
+    # Check for that min_count.
+    with raise_if_dask_computes():
+        actual = getattr(da, func)(dim, skipna=True, min_count=min_count)
+    assert_dask_array(actual, dask)
+    assert_allclose(actual, expected)
+
+    # With min_count being one higher, should get all NaN.
+    min_count += 1
+    expected *= np.nan
+    with raise_if_dask_computes():
+        actual = getattr(da, func)(dim, skipna=True, min_count=min_count)
+    assert_dask_array(actual, dask)
+    assert_allclose(actual, expected)
+
+
 @pytest.mark.parametrize("func", ["sum", "prod"])
 def test_min_count_dataset(func):
     da = construct_dataarray(2, dtype=float, contains_nan=True, dask=False)