Wrap numpy-groupies to speed up Xarray's groupby aggregations

[shoyer]

Is your feature request related to a problem? Please describe.

Xarray's groupby aggregations (e.g., groupby(..).sum()) are very slow compared to pandas, as described in #659.

Describe the solution you'd like

We could speed things up considerably (easily 100x) by wrapping the numpy-groupies package.

Additional context

One challenge is how to handle dask arrays (and other duck arrays). In some cases it might make sense to apply the numpy-groupies function (using apply_ufunc), but in other cases it might be better to stick with the current indexing + concatenate solution. We could either pick some simple heuristics for choosing the algorithm to use on dask arrays, or could just stick with the current algorithm for now.

In particular, it might make sense to stick with the current algorithm if there are a many chunks in the arrays to aggregated along the "grouped" dimension (depending on the size of the unique group values).

[shoyer]

A prototype implementation of the core functionality here can be found in: https://nbviewer.jupyter.org/gist/shoyer/6d6c82bbf383fb717cc8631869678737

[max-sixty]

This looks amazing! Thanks for finding it.

Highly speculative, but would this also be a faster approach to stacking & unstacking? "Form 5 4" in the readme.

[shoyer]

Highly speculative, but would this also be a faster approach to stacking & unstacking? "Form 5 4" in the readme.

I'm not entirely sure, but I suspect something like the approach in #4184 might be more directly relevant for speeding up unstack (at least with NumPy arrays).

[max-sixty]

I'm not entirely sure, but I suspect something like the approach in #4184 might be more directly relevant for speeding up unstack (at least with NumPy arrays).

Great. I need to think through how to do that — the approach of using MultiIndex codes to index the array directly is very elegant — I'll try applying it to stack / unstack as a project.

[max-sixty]

Here's a very quick POC:

from numpy_groupies.aggregate_numba import aggregate

def npg_groupby(da: xr.DataArray, dim, func='sum'):
    group_idx, labels = pd.factorize(da.indexes[dim])
    axis = da.get_axis_num(dim)
    array = npg.aggregate(group_idx=group_idx, a=da, func=func, axis=axis)
    return array

Run on this array:

size_factor = 1000

da = xr.DataArray(
    np.arange(1440 * size_factor).reshape(45 * size_factor, 8, 4),
    dims=("x", "y", "z"),
    coords=dict(x=list(range(45)) * size_factor, y=[1, 2, 3, 4] * 2, z=[1, 2] * 2),
)

It's about 2x as fast, though only generates the numpy array:

%%timeit 
npg_groupby(da, 'x')
# 15 ms ± 130 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%%timeit
da.groupby('x').sum()
# 37.6 ms ± 244 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Any thoughts on any of:

• What's the best way of reconstituting the coords etc, after npg produces the array?
• Presumably we're going to have a fairly different design for this than the existing groupby operations — that design is very nested — wrapping functions and eventually calling .map to loop over each group in python.
• Presumably we're going to need to keep the existing logic around for dask — is it reasonable for an initial version to defer to the existing logic for all dask arrays? (+ @shoyer 's thoughts above on this)

[shoyer]

Something like the resample test case from #4498 might be a good example for finding 100x speed-ups. The main feature of that case is that there are a very large number of groups (only slightly fewer groups than original data points).

[shoyer]

• What's the best way of reconstituting the coords etc, after npg produces the array?

I think we can reuse the existing logic from the _combine method here:

xarray/xarray/core/groupby.py

Line 830 in 97e2625

# assign coord when the applied function does not return that coord

This just gives us an alternative way to calculate applied.

• Presumably we're going to have a fairly different design for this than the existing groupby operations — that design is very nested — wrapping functions and eventually calling .map to loop over each group in python.

Agreed. Hopefully this can live alongside in the GroupBy objects.

• Presumably we're going to need to keep the existing logic around for dask — is it reasonable for an initial version to defer to the existing logic for all dask arrays? (+ @shoyer 's thoughts above on this)

Yes, I agree that we should do this incrementally.