Support shuffle-based groupby aggregations in dask_cudf

[rjzamora]

Description

This PR corresponds to the dask_cudf version of dask/dask#9302 (adding a shuffle-based algorithm for high-cardinality groupby aggregations). The benefits of this algorithm are most significant for cases where split_out>1 is necessary:

agg = ddf.groupby("id").agg({"x": "mean", "y": "max"}, split_out=4, shuffle=True)

NOTES:

• shuffle="explicit-comms" is also supported (when dask_cuda is installed)
• It should be possible to refactor remove some of this code in the future. However, due to some subtle differences between the groupby code in dask.dataframe and dask_cudf, the specialized _shuffle_aggregate is currently necessary.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[wence-]

Ugh, sorry, all of my suggestions have resulted in disaster. Not sure if the better option given we're close to release is go back to the old code and go round again, or to try and fix things up.

I'm inclined perhaps to revert and try again more systematically, WDYT?

[rjzamora]

I'm inclined perhaps to revert and try again more systematically, WDYT?

I think the safest path forward is to stick with the ugly shuffle work-around poposed in this PR (for the 22.10 release), but to rip it out as soon as something like dask/dask#9521 is supported upstream.

What we want for the impending release is to support ddf.shuffle/sort_values/set_index/merge/join(..., shuffle="explicit-comms") and ddf.groupby(...).agg(..., shuffle="explicit-comms"). The current design ended up being uglier than we wanted (and not particularly extensible), but it does "accomplish" this. Therefore, I think we should target a more-extensible redesign for 22.12 (to provide enough time for upstream buy-in and code review).

[rjzamora]

One more thing we can do for 22.10 is to improve the explicit-comms shuffle logic a bit to avoid creating/removing a "_partitions" column when it already exists.

[codecov]

Codecov Report

❗ No coverage uploaded for pull request base (branch-22.10@bcf361f). Click here to learn what that means.
Patch has no changes to coverable lines.

Additional details and impacted files

@@               Coverage Diff               @@
##             branch-22.10   #11800   +/-   ##
===============================================
  Coverage                ?   87.52%           
===============================================
  Files                   ?      133           
  Lines                   ?    21801           
  Branches                ?        0           
===============================================
  Hits                    ?    19081           
  Misses                  ?     2720           
  Partials                ?        0           

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

☔ View full report at Codecov.
📢 Do you have feedback about the report comment? Let us know in this issue.

[rjzamora]

Update: All changes related to explicit-comms have been removed from this PR - The plan for 22.10 is to support shuffle=True or shuffle="tasks" (and to leave out shuffle="explicit-comms").

[wence-]

Thanks, and sorry for all the going around in circles

[shwina]

@gpucibot merge

[rjzamora]

Thanks @wence- and @quasiben for helping to get this in!