Fast GroupBy operations for Dask Arrays

0.11.2 · active · verified Tue Apr 14

Flox is a Python library that provides strategies for fast GroupBy reductions with dask.array, significantly enhancing performance for operations like climatologies, resampling, and histogramming. It was formerly known as `dask_groupby` and integrates seamlessly with xarray to offer more performant GroupBy and Resampling operations.

Warnings

breaking The library was previously known as `dask_groupby`. Code relying on the old package name or import paths will break.
Fix: Update imports from `dask_groupby` to `flox` and adjust any related API calls.
gotcha When `flox` (version >= 2022.06.0) is installed, Xarray will automatically use `flox` by default for its `.groupby`, `.groupby_bins`, and `.resample` operations. This implicit usage can change performance characteristics or expose underlying `flox` issues.
Fix: Be aware that Xarray's GroupBy methods might be leveraging `flox`. To debug performance or unexpected behavior, consider isolating `flox` or temporarily uninstalling it to revert to Xarray's default GroupBy engine.
gotcha Custom reductions specified using `Aggregation` instances might not be fully functional or have undefined behavior in certain scenarios.
Fix: Rely on the built-in reduction functions (e.g., 'mean', 'sum') first. If custom logic is required, thoroughly test its behavior and consult documentation/issues for known limitations.
gotcha High memory usage can occur with `flox` aggregations in Dask, particularly when lower-level tasks (e.g., data loading) continue running while higher-level reduction tasks are uncomputed.
Fix: Monitor Dask dashboard for task graph progression. Consider rechunking strategies, explicit method selection (e.g., `method='map-reduce'`), or breaking down computations into smaller steps if memory issues persist.
gotcha For Dask arrays, `flox` uses heuristics (since v0.9.0) to choose the optimal parallel algorithm (`map-reduce`, `blockwise`, `cohorts`). While generally robust, specific data distributions or chunking patterns might benefit from explicitly setting the `method` parameter in `groupby_reduce` or `xarray_reduce`.
Fix: If performance is not as expected, experiment with `method` and `reindex` arguments (e.g., `method='map-reduce'`, `method='blockwise'`, `method='cohorts'`) to find the best strategy for your data.

Install

pip install flox Install stable release

Imports

groupby_reduce
```
from flox import groupby_reduce
```
xarray_reduce
```
from flox.xarray import xarray_reduce
```
For Xarray objects and lazy grouping by Dask arrays.

Quickstart

This quickstart demonstrates how to use `flox.groupby_reduce` with a Dask array and a NumPy array of group labels to compute the mean for each group. The `expected_groups` argument ensures all groups are present in the output, even if some are empty.

import dask.array as da
from flox import groupby_reduce
import numpy as np

# Create a sample Dask array
data = da.random.random((1000, 10), chunks=(100, 10))

# Create a 'by' array for grouping (e.g., categories 0-9)
groups = np.random.randint(0, 10, size=1000)

# Perform a GroupBy reduction (e.g., mean)
result_mean, group_labels = groupby_reduce(
    data, groups, func="mean", expected_groups=np.arange(10)
)

print("Grouped Means (first 5 groups):\n", result_mean.compute()[:5])
print("Group Labels:\n", group_labels)

view raw JSON →