Use numbagg for rolling methods

xarray/core/rolling.py

@@ -14,7 +14,7 @@
 from xarray.core.options import OPTIONS, _get_keep_attrs
 from xarray.core.pycompat import is_duck_dask_array
 from xarray.core.types import CoarsenBoundaryOptions, SideOptions, T_Xarray
-from xarray.core.utils import either_dict_or_kwargs
+from xarray.core.utils import either_dict_or_kwargs, module_available
 
 try:
     import bottleneck
@@ -153,14 +153,21 @@ def _reduce_method(  # type: ignore[misc]
             array_agg_func = getattr(duck_array_ops, name)
 
         bottleneck_move_func = getattr(bottleneck, "move_" + name, None)
+        if module_available("numbagg"):
+            import numbagg
+
+            numbagg_move_func = getattr(numbagg, "move_" + name, None)
+        else:
+            numbagg_move_func = None
 
         def method(self, keep_attrs=None, **kwargs):
             keep_attrs = self._get_keep_attrs(keep_attrs)
 
             return self._numpy_or_bottleneck_reduce(
-                array_agg_func,
-                bottleneck_move_func,
-                rolling_agg_func,
+                array_agg_func=array_agg_func,
+                bottleneck_move_func=bottleneck_move_func,
+                numbagg_move_func=numbagg_move_func,
+                rolling_agg_func=rolling_agg_func,
                 keep_attrs=keep_attrs,
                 fillna=fillna,
                 **kwargs,
@@ -510,9 +517,47 @@ def _counts(self, keep_attrs: bool | None) -> DataArray:
         )
         return counts
 
-    def _bottleneck_reduce(self, func, keep_attrs, **kwargs):
-        from xarray.core.dataarray import DataArray
+    def _numbagg_reduce(self, func, keep_attrs, **kwargs):
+        # Some of this is copied from `_bottleneck_reduce`, we could reduce this as part
+        # of a wider refactor.
+
+        axis = self.obj.get_axis_num(self.dim[0])
+
+        padded = self.obj.variable
+        if self.center[0]:
+            if is_duck_dask_array(padded.data):
+                # workaround to make the padded chunk size larger than
+                # self.window - 1
+                shift = -(self.window[0] + 1) // 2
+                offset = (self.window[0] - 1) // 2
+                valid = (slice(None),) * axis + (
+                    slice(offset, offset + self.obj.shape[axis]),
+                )
+            else:
+                shift = (-self.window[0] // 2) + 1
+                valid = (slice(None),) * axis + (slice(-shift, None),)
+            padded = padded.pad({self.dim[0]: (0, -shift)}, mode="constant")
+
+        if is_duck_dask_array(padded.data):
+            raise AssertionError("should not be reachable")
+        else:
+            values = func(
+                padded.data,
+                window=self.window[0],
+                min_count=self.min_periods,
+                axis=axis,
+            )
+
+        if self.center[0]:
+            values = values[valid]
 
+        attrs = self.obj.attrs if keep_attrs else {}
+
+        return self.obj.__class__(
+            values, self.obj.coords, attrs=attrs, name=self.obj.name
+        )
+
+    def _bottleneck_reduce(self, func, keep_attrs, **kwargs):
         # bottleneck doesn't allow min_count to be 0, although it should
         # work the same as if min_count = 1
         # Note bottleneck only works with 1d-rolling.
@@ -550,12 +595,15 @@ def _bottleneck_reduce(self, func, keep_attrs, **kwargs):
 
         attrs = self.obj.attrs if keep_attrs else {}
 
-        return DataArray(values, self.obj.coords, attrs=attrs, name=self.obj.name)
+        return self.obj.__class__(
+            values, self.obj.coords, attrs=attrs, name=self.obj.name
+        )
 
     def _numpy_or_bottleneck_reduce(
         self,
         array_agg_func,
         bottleneck_move_func,
+        numbagg_move_func,
         rolling_agg_func,
         keep_attrs,
         fillna,
@@ -571,6 +619,26 @@ def _numpy_or_bottleneck_reduce(
             )
             del kwargs["dim"]
 
+        if (
+            OPTIONS["use_numbagg"]
+            and numbagg_move_func is not None
+            # TODO: can we allow this with numbagg?
+            and not is_duck_dask_array(self.obj.data)
+            and self.obj.data.dtype.kind not in "Ob"
+        ):
+            import numbagg
+
+            # Numbagg has a default ddof of 1. I (@max-sixty) think we should make
+            # this the default in xarray too, but until we do, don't use numbagg for
+            # std and var unless ddof is set to 1.
+            if (
+                numbagg_move_func not in [numbagg.move_std, numbagg.move_var]
+                or kwargs.get("ddof") == 1
+            ):
+                return self._numbagg_reduce(
+                    numbagg_move_func, keep_attrs=keep_attrs, **kwargs
+                )
+
         if (
             OPTIONS["use_bottleneck"]
             and bottleneck_move_func is not None

xarray/tests/test_rolling.py

@@ -30,7 +30,7 @@
     @pytest.mark.parametrize("size", [1, 2, 3, 7])
     def test_rolling_iter(self, da: DataArray, center: bool, size: int) -> None:
         rolling_obj = da.rolling(time=size, center=center)
         rolling_obj_mean = rolling_obj.mean()
 
         assert len(rolling_obj.window_labels) == len(da["time"])
         assert_identical(rolling_obj.window_labels, da["time"])
@@ -98,7 +98,9 @@
         expected = getattr(bn, func_name)(
             da.values, window=7, axis=1, min_count=min_periods
         )
-        assert_array_equal(actual.values, expected)
+
+        # Using assert_allclose because we get tiny (1e-17) differences in numbagg.
+        np.testing.assert_allclose(actual.values, expected)
 
         with pytest.warns(DeprecationWarning, match="Reductions are applied"):
             getattr(rolling_obj, name)(dim="time")
@@ -280,28 +282,33 @@
     @pytest.mark.parametrize("min_periods", (None, 1))
     @pytest.mark.parametrize("name", ("sum", "mean", "max"))
     def test_ndrolling_reduce(self, da, center, min_periods, name) -> None:
-        rolling_obj = da.rolling(time=3, x=2, center=center, min_periods=min_periods)
-
-        actual = getattr(rolling_obj, name)()
-        expected = getattr(
-            getattr(
-                da.rolling(time=3, center=center, min_periods=min_periods), name
-            )().rolling(x=2, center=center, min_periods=min_periods),
-            name,
-        )()
-
-        assert_allclose(actual, expected)
-        assert actual.dims == expected.dims
-
-        if name in ["mean"]:
-            # test our reimplementation of nanmean using np.nanmean
-            expected = getattr(rolling_obj.construct({"time": "tw", "x": "xw"}), name)(
-                ["tw", "xw"]
+        # FIXME
+        # with xr.set_options(use_numbagg=False):
+        if True:
+            rolling_obj = da.rolling(
+                time=3, x=2, center=center, min_periods=min_periods
             )
-            count = rolling_obj.count()
-            if min_periods is None:
-                min_periods = 1
-            assert_allclose(actual, expected.where(count >= min_periods))
+
+            actual = getattr(rolling_obj, name)()
+            expected = getattr(
+                getattr(
+                    da.rolling(time=3, center=center, min_periods=min_periods), name
+                )().rolling(x=2, center=center, min_periods=min_periods),
+                name,
+            )()
+
+            assert_allclose(actual, expected)
+            assert actual.dims == expected.dims
+
+            if name in ["mean"]:
+                # test our reimplementation of nanmean using np.nanmean
+                expected = getattr(
+                    rolling_obj.construct({"time": "tw", "x": "xw"}), name
+                )(["tw", "xw"])
+                count = rolling_obj.count()
+                if min_periods is None:
+                    min_periods = 1
+                assert_allclose(actual, expected.where(count >= min_periods))
 
     @pytest.mark.parametrize("center", (True, False, (True, False)))
     @pytest.mark.parametrize("fill_value", (np.nan, 0.0))
@@ -577,7 +584,7 @@
             )
         else:
             raise ValueError
         assert_array_equal(actual[key].values, expected)
 
         # Test center
         rolling_obj = ds.rolling(time=7, center=center)
@@ -726,7 +733,7 @@
             )().rolling(x=3, center=center, min_periods=min_periods),
             name,
         )()
         assert_allclose(actual, expected)
         assert actual.dims == expected.dims
 
         # Do it in the opposite order