Repository navigation
Use flox for median in groupby/resample operations - #11239
Conversation
Add flox support to median() methods in: - DataArrayGroupByAggregations - DatasetGroupByAggregations - DataArrayResampleAggregations - DatasetResampleAggregations This aligns the implementation with the documentation which already claimed flox was used when available. The fix provides significant performance improvements when flox can process the data. A fallback to the non-flox implementation is included for cases where flox's median aggregation requires blockwise processing but the data chunking doesn't support it. Closes pydata#11238
There was a problem hiding this comment.
Thanks for taking this on!
Let's raise Now that I think about it; we can do this in flox automatically: xarray-contrib/flox#501FutureWarning here asking the user to shuffle_to_chunks or rechunk appropriately and that this behaviour will be deprecated in 6 months (please open an issue for that).
The reason I hadn't done this so far is that dask will do some auto-rechunking to make it work; and so some workloads will break.
Also this file is generated by generate_aggregations.py. Please make the edit there instead
themavik
left a comment
themavik
left a comment
There was a problem hiding this comment.
this could break silently if flox rewords the ValueError for non-blockwise median; fallback keys off an exact substring.
Will dot! Should I also run |
Yes, and hmmm.... the output shouldn't chnage that much. |
…kwise" Edit generate_aggregations.py instead of the generated _aggregations.py file and pass method="blockwise" to flox, which lets flox automatically rechunk 1-D groupers so the call works regardless of input chunking. The previous try/except fallback keyed on a flox error message substring, which could break silently if flox ever reworded that error.
github-actions
Bot
added
topic-NamedArray
0ebee6f to
71b4b96
Compare
github-actions
Bot
added
topic-backends
topic-indexing
topic-groupby
topic-dask
topic-plotting
topic-testing
topic-performance
topic-cftime
topic-zarr
Resolves merge conflicts in: - xarray/util/generate_aggregations.py: combine flox_call_kwargs (median dispatch) with upstream's aggregation_type/long_name support - xarray/core/_aggregations.py: regenerate via the merged generator
After regenerating _aggregations.py files with the merged generator, the doctest expectations no longer matched the actual output. Running pytest with --accept updated the docstrings to include the expected output so the doctests pass.
During the merge of upstream/main into median-flox, the
method_is_not_flox_supported tuple in generate_aggregations.py ended up as
("cumsum", "cumprod"). That meant cumsum fell through to the non-flox path
and broke under dask + lazy groupby (macos-latest | test-py314 and
ubuntu-latest | test-py314-no-numba). It also kept median in the list,
defeating the goal of this PR.
The correct tuple is just ("cumprod",) since flox doesn't implement it,
but does implement median (this PR) and cumsum (since flox 0.10.5).
Regenerated core/_aggregations.py via the merged generator and ran
pytest --doctest-modules --accept to update docstrings.
The previous fix used method.name in ("cumprod",) which ruff flagged as a
membership test against a single-item container. Use a direct equality
check instead.
Regenerated core/_aggregations.py.
|
Finally found time to come back to this PR, sorry for the delay! It should be good to review again. I may have broken the labels bot by force-pushing somehow, sorry about that! |
Use kwargs.setdefault so callers can still override flox method while median defaults to blockwise.
* Update whats-new for v2026.09.0 release - Rename the unreleased section to v2026.09.0 and add the release summary and contributor list - Move entries for #11348, #11292, #11290 and #11239 that were merged into older release sections - Add missing entries for #11552, #11521, #11486 and #11513 - Add missing PR links, fix the #11547 link and other markup - Fix the documented default of the use_bottleneck option - Add contributor names to the typos allow-list * Mention plot.lines and last Python 3.11 release in whats-new
Since pydata#11239, groupby median on dask arrays uses flox with method="blockwise". flox can only rechunk so that chunk boundaries line up with group boundaries if the members of each group are contiguous. With interleaved labels spread over multiple chunks this raised an InvalidIndexError. Shuffle every group into a single chunk first. This can be restricted to older flox versions once xarray-contrib/flox#501 is fixed. Fixes pydata#11651 Co-authored-by: Claude <[email protected]>
Add flox support to median() methods in:
This aligns the implementation with the documentation which already claimed flox was used when available. The fix provides significant performance improvements when flox can process the data. A fallback to the non-flox implementation is included for cases where flox's median aggregation requires blockwise processing but the data chunking doesn't support it.
Closes #11238
whats-new.rst