Explore moving functions from reduce to apply_ufunc

Author: max-sixty

xarray/core/rolling_exp.py

@@ -5,6 +5,7 @@
 
 import numpy as np
 
+from xarray.core.computation import apply_ufunc
 from xarray.core.options import _get_keep_attrs
 from xarray.core.pdcompat import count_not_none
 from xarray.core.pycompat import is_duck_dask_array
@@ -128,9 +129,18 @@ def mean(self, keep_attrs: bool | None = None) -> T_DataWithCoords:
         if keep_attrs is None:
             keep_attrs = _get_keep_attrs(default=True)
 
-        return self.obj.reduce(
-            move_exp_nanmean, dim=self.dim, alpha=self.alpha, keep_attrs=keep_attrs
-        )
+        dim_order = self.obj.dims
+
+        return apply_ufunc(
+            move_exp_nanmean,
+            self.obj,
+            input_core_dims=[[self.dim]],
+            kwargs=dict(alpha=self.alpha, axis=-1),
+            output_core_dims=[[self.dim]],
+            exclude_dims={self.dim},
+            keep_attrs=keep_attrs,
+            on_missing_core_dim="copy",
+        ).transpose(*dim_order)
 
     def sum(self, keep_attrs: bool | None = None) -> T_DataWithCoords:
         """
@@ -155,6 +165,15 @@ def sum(self, keep_attrs: bool | None = None) -> T_DataWithCoords:
         if keep_attrs is None:
             keep_attrs = _get_keep_attrs(default=True)
 
-        return self.obj.reduce(
-            move_exp_nansum, dim=self.dim, alpha=self.alpha, keep_attrs=keep_attrs
-        )
+        dim_order = self.obj.dims
+
+        return apply_ufunc(
+            move_exp_nansum,
+            self.obj,
+            input_core_dims=[[self.dim]],
+            kwargs=dict(alpha=self.alpha, axis=-1),
+            output_core_dims=[[self.dim]],
+            exclude_dims={self.dim},
+            keep_attrs=keep_attrs,
+            on_missing_core_dim="copy",
+        ).transpose(*dim_order)