Use cumsum from flox

xarray/core/_aggregations.py

@@ -6647,6 +6647,13 @@ def _flox_reduce(
     ) -> DataArray:
         raise NotImplementedError()
 
+    def _flox_scan(
+        self,
+        dim: Dims,
+        **kwargs: Any,
+    ) -> DataArray:
+        raise NotImplementedError()
+
     def count(
         self,
         dim: Dims = None,
@@ -7904,13 +7911,27 @@ def cumsum(
           * time     (time) datetime64[ns] 48B 2001-01-31 2001-02-28 ... 2001-06-30
             labels   (time) <U1 24B 'a' 'b' 'c' 'c' 'b' 'a'
         """
-        return self.reduce(
-            duck_array_ops.cumsum,
-            dim=dim,
-            skipna=skipna,
-            keep_attrs=keep_attrs,
-            **kwargs,
-        )
+        if (
+            flox_available
+            and OPTIONS["use_flox"]
+            and contains_only_chunked_or_numpy(self._obj)
+        ):
+            return self._flox_scan(
+                func="cumsum",
+                dim=dim,
+                skipna=skipna,
+                # fill_value=fill_value,
+                keep_attrs=keep_attrs,
+                **kwargs,
+            )
+        else:
+            return self.reduce(
+                duck_array_ops.cumsum,
+                dim=dim,
+                skipna=skipna,
+                keep_attrs=keep_attrs,
+                **kwargs,
+            )
 
     def cumprod(
         self,

xarray/core/groupby.py

@@ -13,6 +13,7 @@
 from packaging.version import Version
 
 from xarray.computation import ops
+from xarray.computation.apply_ufunc import apply_ufunc
 from xarray.computation.arithmetic import (
     DataArrayGroupbyArithmetic,
     DatasetGroupbyArithmetic,
@@ -1028,6 +1029,26 @@ def _maybe_unstack(self, obj):
 
         return obj
 
+    def _parse_dim(self, dim: Dims) -> tuple[Hashable, ...]:
+        parsed_dim: tuple[Hashable, ...]
+        if isinstance(dim, str):
+            parsed_dim = (dim,)
+        elif dim is None:
+            parsed_dim_list = list()
+            # preserve order
+            for dim_ in itertools.chain(
+                *(grouper.codes.dims for grouper in self.groupers)
+            ):
+                if dim_ not in parsed_dim_list:
+                    parsed_dim_list.append(dim_)
+            parsed_dim = tuple(parsed_dim_list)
+        elif dim is ...:
+            parsed_dim = tuple(self._original_obj.dims)
+        else:
+            parsed_dim = tuple(dim)
+
+        return parsed_dim
+
     def _flox_reduce(
         self,
         dim: Dims,
@@ -1088,22 +1109,7 @@ def _flox_reduce(
                 # set explicitly to avoid unnecessarily accumulating count
                 kwargs["min_count"] = 0
 
-        parsed_dim: tuple[Hashable, ...]
-        if isinstance(dim, str):
-            parsed_dim = (dim,)
-        elif dim is None:
-            parsed_dim_list = list()
-            # preserve order
-            for dim_ in itertools.chain(
-                *(grouper.codes.dims for grouper in self.groupers)
-            ):
-                if dim_ not in parsed_dim_list:
-                    parsed_dim_list.append(dim_)
-            parsed_dim = tuple(parsed_dim_list)
-        elif dim is ...:
-            parsed_dim = tuple(obj.dims)
-        else:
-            parsed_dim = tuple(dim)
+        parsed_dim = self._parse_dim(dim)
 
         # Do this so we raise the same error message whether flox is present or not.
         # Better to control it here than in flox.
@@ -1202,6 +1208,85 @@ def _flox_reduce(
 
         return result
 
+    def _flox_scan(
+        self,
+        dim: Dims,
+        *,
+        func: str,
+        keep_attrs: bool | None = None,
+        skipna: bool | None = None,
+        **kwargs: Any,
+    ) -> DataArray:
+        from flox import groupby_scan
+
+        obj = self._original_obj
+
+        if skipna or (
+            skipna is None and isinstance(func, str) and obj.dtype.kind in "cfO"
+        ):
+            if "nan" not in func and func not in ["all", "any", "count"]:
+                func = f"nan{func}"
+
+        # if keep_attrs is None:
+        #     keep_attrs = _get_keep_attrs(default=True)
+
+        parsed_dim = self._parse_dim(dim)
+
+        axis_ = obj.get_axis_num(parsed_dim)
+        axis = (axis_,) if isinstance(axis_, int) else axis_
+        codes = tuple(g.codes for g in self.groupers)
+        # g = groupby_scan(
+        #     obj.data,
+        #     *codes,
+        #     func=func,
+        #     expected_groups=None,
+        #     axis=axis,
+        #     dtype=None,
+        #     method=None,
+        #     engine=None,
+        # )
+        # result = obj.copy(data=g)
+
+        # return result
+
+        actual = apply_ufunc(
+            groupby_scan,
+            obj,
+            *codes,
+            # input_core_dims=input_core_dims,
+            # for xarray's test_groupby_duplicate_coordinate_labels
+            # exclude_dims=set(dim_tuple),
+            # output_core_dims=[output_core_dims],
+            dask="allowed",
+            # dask_gufunc_kwargs=dict(
+            #     output_sizes=output_sizes,
+            #     output_dtypes=[dtype] if dtype is not None else None,
+            # ),
+            keep_attrs=(
+                _get_keep_attrs(default=True) if keep_attrs is None else keep_attrs
+            ),
+            kwargs=dict(
+                func=func,
+                expected_groups=None,
+                axis=axis,
+                dtype=None,
+                method=None,
+                engine=None,
+            ),
+        )
+
+        return actual
+
+        # xarray_reduce(
+        #     obj.drop_vars(non_numeric.keys()),
+        #     *codes,
+        #     dim=parsed_dim,
+        #     expected_groups=expected_groups,
+        #     isbin=False,
+        #     keep_attrs=keep_attrs,
+        #     **kwargs,
+        # )
+
     def fillna(self, value: Any) -> T_Xarray:
         """Fill missing values in this object by group.