Explore flox for improving Xarray's groupby() performance

[tomvothecoder]

Is your feature request related to a problem?

According to Xarray's documentation, groupby() and resample() triggers computation on all blocks which means it loads all data into memory and does not operate on them lazily. This results in a costly operation.

Do your spatial and temporal indexing (e.g. .sel() or .isel()) early in the pipeline, especially before calling resample() or groupby(). Grouping and resampling triggers some computation on all the blocks, which in theory should commute with indexing, but this optimization hasn’t been implemented in Dask yet. (See Dask issue #746).

— https://docs.xarray.dev/en/stable/user-guide/dask.html#optimization-tips

This is where flox comes along:

More generally, groupby() is a costly operation and will perform a lot better if the flox package is installed. See the flox documentation for more. By default Xarray will use flox if installed.

— https://docs.xarray.dev/en/stable/user-guide/dask.html#optimization-tips

flox mainly provides strategies for fast GroupBy reductions with dask.array. flox uses the MapReduce paradigm (or a “tree reduction”) to run the GroupBy operation in a parallel-native way totally avoiding a sort or shuffle operation.

— https://flox.readthedocs.io/en/latest/

Describe the solution you'd like

• Capture performance metrics using datasets of varying sizes to see the impact of flox with xCDAT APIs

If performance increases are substantial:

• Include flox as an xCDAT dependency
  • Examples: https://flox.readthedocs.io/en/latest/user-stories/climatology.html
• Specify explicit reduction method where groupby() calls are made internally

Describe alternatives you've considered

• Add information on our docs about how users can use flox to speed up groupby operations, including spatial and temporal averaging.

Additional context

One challenge might be that we need to constraint the minimum version of xarray to >=2022.6.0 because that is when flox was first implemented in xarray.