Use cumsum from flox

[Illviljan]

• Closes cumsum drops index coordinates #6528
• Tests added
• User visible changes (including notable bug fixes) are documented in whats-new.rst
• New functions/methods are listed in api.rst

The non-flox version reduces chunksizes significantly:

x = xr.DataArray([1, 1, 1, 1, 1], name="x").chunk()
grp_idx = xr.DataArray([-1, 0, 0, -1, 1])
with xr.set_options(use_flox=False):
    print(x.groupby(grp_idx).cumsum())
<xarray.DataArray 'x' (dim_0: 5)> Size: 40B
dask.array<getitem, shape=(5,), dtype=int64, chunksize=(2,), chunktype=numpy.ndarray>
Dimensions without coordinates: dim_0

With flox the chunksize is retained:

x = xr.DataArray([1, 1, 1, 1, 1], name="x").chunk()
grp_idx = xr.DataArray([-1, 0, 0, -1, 1])
with xr.set_options(use_flox=True):
    print(x.groupby(grp_idx).cumsum())
<xarray.DataArray 'x' (dim_0: 5)> Size: 40B
dask.array<_finalize_scan, shape=(5,), dtype=int64, chunksize=(5,), chunktype=numpy.ndarray>
Dimensions without coordinates: dim_0

[dcherian]

FYI in a future PR, I'd like to use the new flox.is_supported_aggregation here. It's a little smarter about this dispatching. We'll also have to figure out what to do about median which currently auto-rechunks so it always works.

[dcherian]

is this needed?

[Illviljan]

I thought this looked nice and it was clear there was tweak to fix cumsum/cumprod:

            out = self.reduce(
                duck_array_ops.cumsum,
                dim=dim,
                skipna=skipna,
                keep_attrs=keep_attrs,
                **kwargs,
            )
            return out.assign_coords(self._obj.coords)

Then for consistency and readability I followed that pattern on the others.

[dcherian]

Can we apply the assign_coords fix for Dataset.cumsum et al too?

[dcherian]

The use_dask here should mean grouping a dask array by a numpy array. That will work always.