flox 0.2.0

GroupBy operations for dask.array

flox

This project explores strategies for fast GroupBy reductions with dask.array. It used to be called dask_groupby

This repo explores strategies for a distributed GroupBy with dask arrays. It was motivated by

1. Dask Dataframe GroupBy blogpost
2. numpy_groupies in Xarray issue

(See a presentation about this package).

Acknowledgements

This work was funded in part by NASA-ACCESS 80NSSC18M0156 "Community tools for analysis of NASA Earth Observing System Data in the Cloud" (PI J. Hamman), and NCAR's Earth System Data Science Initiative. It was motivated by many discussions in the Pangeo community.

API

There are three functions

1. flox.groupby_reduce(dask_array, by_dask_array, "mean") "pure" dask array interface
2. flox.xarray.xarray_reduce(xarray_object, by_dataarray, "mean") "pure" xarray interface

Implementation

The core GroupBy operation is outsourced to numpy_groupies. The GroupBy reduction is first applied blockwise. Those intermediate results are combined by concatenating to form a new array which is then reduced again. The combining of intermediate results uses dask's _tree_reduce till all group results are in one block. At that point the result is "finalized" and returned to the user. Here is an example of writing a custom Aggregation (again inspired by dask.dataframe)

    mean = Aggregation(
        # name used for dask tasks
        name="mean",
        # blockwise reduction
        chunk=("sum", "count"),
        # combine intermediate results: sum the sums, sum the counts
        combine=("sum", "sum"),
        # generate final result as sum / count
        finalize=lambda sum_, count: sum_ / count,
        # Used when "reindexing" at combine-time
        fill_value=0,
    )

Using _tree_reduce complicates the implementation. An alternative simpler implementation would be to use the "tensordot" trick. But this requires knowledge of "expected group labels" at compute-time.