Use cumsum from flox #10987

Illviljan · 2025-12-06T13:44:27Z

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

for more information, see https://pre-commit.ci

…o cumsum_flox

for more information, see https://pre-commit.ci

…o cumsum_flox

for more information, see https://pre-commit.ci

xarray/core/groupby.py

…o cumsum_flox

for more information, see https://pre-commit.ci

…o cumsum_flox

for more information, see https://pre-commit.ci

…o cumsum_flox

for more information, see https://pre-commit.ci

xarray/core/groupby.py

Co-authored-by: Deepak Cherian <[email protected]>

for more information, see https://pre-commit.ci

xarray/tests/test_groupby.py

for more information, see https://pre-commit.ci

…o cumsum_flox

dcherian · 2025-12-15T21:00:58Z

xarray/core/datatree.py

-                **kwargs,
-            )
+
+            # Prefer Dataset.func(...) over Dataset.reduce(func, ...):


why is this needed

cc @TomNicholas

Because of #6528 and #10987 (comment). Since the problem of Datasets dropping the coords is solved in .cumsum in this PR, DataTree must also use ds.cumsum instead of ds.reduce("cumsum") in order to avoid dropping the coordinates.

dcherian · 2025-12-16T02:55:57Z

xarray/tests/test_dataset.py

    )
    @pytest.mark.parametrize("func", ["cumsum", "cumprod"])
-    def test_reduce_cumsum_test_dims(self, reduct, expected, func) -> None:
+    def test_reduce_cumsum_test_dims(self, reduct, func) -> None:


can we also assert that coordinates & indexes are preserved please

dcherian · 2025-12-16T02:56:17Z

xarray/tests/test_groupby.py

+    with xr.set_options(use_flox=use_flox):
+        if use_dask:
+            ds = ds.chunk()
+            if use_lazy_group_idx:


this is always False

use cumsum from flox

776bc5a

github-actions bot added the topic-groupby label Dec 6, 2025

pre-commit-ci bot and others added 13 commits December 6, 2025 13:44

[pre-commit.ci] auto fixes from pre-commit.com hooks

ae27632

for more information, see https://pre-commit.ci

Update groupby.py

a5f9326

Update groupby.py

50ccca4

[pre-commit.ci] auto fixes from pre-commit.com hooks

f55531e

for more information, see https://pre-commit.ci

Update groupby.py

06ac372

Merge branch 'cumsum_flox' of https://github.com/Illviljan/xarray int…

31244e6

…o cumsum_flox

Update groupby.py

dd47536

[pre-commit.ci] auto fixes from pre-commit.com hooks

e867f12

for more information, see https://pre-commit.ci

Update groupby.py

88e0ebc

[pre-commit.ci] auto fixes from pre-commit.com hooks

181d4a3

for more information, see https://pre-commit.ci

use apply_ufunc for dataset and dataarray handling

a82ec39

Merge branch 'cumsum_flox' of https://github.com/Illviljan/xarray int…

6c6abed

…o cumsum_flox

[pre-commit.ci] auto fixes from pre-commit.com hooks

24c3f1d

for more information, see https://pre-commit.ci

dcherian reviewed Dec 6, 2025

View reviewed changes

xarray/core/groupby.py Show resolved Hide resolved

dcherian reviewed Dec 6, 2025

View reviewed changes

xarray/core/groupby.py Show resolved Hide resolved

Illviljan and others added 11 commits December 6, 2025 16:21

Update groupby.py

d8d0eaa

Merge branch 'cumsum_flox' of https://github.com/Illviljan/xarray int…

55ff46a

…o cumsum_flox

[pre-commit.ci] auto fixes from pre-commit.com hooks

33d1360

for more information, see https://pre-commit.ci

sync protocols with each other

c97ae98

Merge branch 'cumsum_flox' of https://github.com/Illviljan/xarray int…

06b52ae

…o cumsum_flox

typing

84f9b44

[pre-commit.ci] auto fixes from pre-commit.com hooks

2978877

for more information, see https://pre-commit.ci

add dataset and version requirement

0a9adee

Merge branch 'cumsum_flox' of https://github.com/Illviljan/xarray int…

ae9a3d8

…o cumsum_flox

[pre-commit.ci] auto fixes from pre-commit.com hooks

c056d1f

for more information, see https://pre-commit.ci

Update _aggregations.py

d4873b9

dcherian reviewed Dec 6, 2025

View reviewed changes

xarray/core/groupby.py Outdated Show resolved Hide resolved

Update xarray/core/groupby.py

21cbde2

Co-authored-by: Deepak Cherian <[email protected]>

Illviljan and others added 5 commits December 12, 2025 21:14

Update test_groupby.py

1927923

[pre-commit.ci] auto fixes from pre-commit.com hooks

45e9423

for more information, see https://pre-commit.ci

Update generate_aggregations.py

3a9cee8

Update _aggregations.py

a3281e6

Update _aggregations.py

f29cdd0

dcherian reviewed Dec 12, 2025

View reviewed changes

xarray/tests/test_groupby.py Show resolved Hide resolved

Illviljan added 5 commits December 13, 2025 12:44

Update test_groupby.py

76349a2

Update generate_aggregations.py

52df016

Update _aggregations.py

5b88b81

Update _aggregations.py

e43220e

Datatree should use Dataset.func rather than Dataset.reduce(func)

fec8438

github-actions bot added the topic-DataTree Related to the implementation of a DataTree class label Dec 13, 2025

Illviljan and others added 14 commits December 13, 2025 14:05

docstsrings for datatree can now be fixed

9f971c2

Update _aggregations.py

44cebc3

Update test_dataset.py

c16fc9e

Update test_dataset.py

defcc44

Merge branch 'main' into cumsum_flox

7b5be19

Update generate_aggregations.py

d08804f

Update _aggregations.py

35aeb3b

Update generate_aggregations.py

5ebce3c

Update generate_aggregations.py

4aad0dc

Update _aggregations.py

eff669e

[pre-commit.ci] auto fixes from pre-commit.com hooks

ec97813

for more information, see https://pre-commit.ci

mapping in collections

362af7f

Merge branch 'cumsum_flox' of https://github.com/Illviljan/xarray int…

fad2d3f

…o cumsum_flox

Merge branch 'main' into cumsum_flox

4df74a9

dcherian reviewed Dec 15, 2025

View reviewed changes

add coords

00f9d18

dcherian reviewed Dec 16, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

Use cumsum from flox #10987

Use cumsum from flox #10987

Illviljan commented Dec 6, 2025 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dcherian Dec 15, 2025

Uh oh!

Illviljan Dec 15, 2025 •

edited

Loading

Uh oh!

dcherian Dec 16, 2025

Uh oh!

dcherian Dec 16, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Use cumsum from flox #10987

Are you sure you want to change the base?

Use cumsum from flox #10987

Conversation

Illviljan commented Dec 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dcherian Dec 15, 2025

Choose a reason for hiding this comment

Uh oh!

Illviljan Dec 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dcherian Dec 16, 2025

Choose a reason for hiding this comment

Uh oh!

dcherian Dec 16, 2025

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Illviljan commented Dec 6, 2025 •

edited

Loading

Illviljan Dec 15, 2025 •

edited

Loading