WIP: Dask array expressions#508
Open
mrocklin wants to merge 6 commits into
Open
Conversation
Use dask-array expressions for flox groupby operations and route standalone dask_array inputs through an isolated backend. Keep the legacy dask.array path side-effect free, explicitly reject mixed dask backends, and cover xarray, unknown groups, cohorts, and import isolation in focused tests. Co-Authored-By: Codex <codex@openai.com>
Parametrize selected flox tests over legacy dask.array and standalone dask_array while avoiding collection-time dask_array imports that alter xarray's chunk-manager registry. Add a standalone arg-reduction preprocess path so dask_array stays on expression graphs for arg reductions. Co-Authored-By: Codex <codex@openai.com>
for more information, see https://pre-commit.ci
Make optional dask-array typing explicit for mypy, skip the legacy-Dask subprocess test when Dask is absent, and include rich in docs dependencies for the cubed notebook. Co-Authored-By: Codex <codex@openai.com>
Dask-array rolling-window reductions can simplify to a different block grid than the public pre-lowered expression reports. Flox was planning groupby chunks from the pre-simplified grid, which made the final dask-array blockwise extraction inconsistent after lowering and broke map-reduce, cohorts, and xarray's default flox path. Explicit blockwise had the same issue one step earlier because rechunk_for_blockwise preserved stale spatial chunking before the dask-array backend saw the array. Simplify standalone dask-array inputs before chunk unification/planning, and before blockwise group-boundary rechunking. Add a small rolling groupby regression that checks all flox paths compute correctly and keep the expected coarser spatial chunk grid. Co-Authored-By: Codex <codex@openai.com>
dcherian
reviewed
Jun 29, 2026
| aggregate = partial(_aggregate, combine=combine, agg=agg, fill_value=fill_value, reindex=reindex) | ||
|
|
||
| if method == "map-reduce": | ||
| reduced = _tree_reduce( |
Collaborator
There was a problem hiding this comment.
use dask_array.reduction
dcherian
reviewed
Jun 29, 2026
| group_chunks = ((len(expected_groups),),) | ||
|
|
||
| else: | ||
| assert chunks_cohorts |
Collaborator
There was a problem hiding this comment.
iterate over cohort;s use reduction for tree-reduction over each cohort; then concatenate across cohorts.
Replace the private dask_array _tree_reduce import in the standalone dask-array backend with the public dask_array.reduction API for both map-reduce and per-cohort reductions. Dask-array simplification can change the block grid used by rolling-window inputs, so recompute cohort block membership after simplifying and unifying chunks. That keeps cohort subsetting aligned with the graph that will actually be lowered and computed. Co-Authored-By: Codex <codex@openai.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds
dask-arrayexpression support to flox.Marking this as WIP mostly because I haven't reviewed it yet. I didn't want to forget about it though.