[DISCUSS] Improve client/scheduler performance during shuffling

[rjzamora]

Lets use this issue to coordinate some ongoing efforts to improve client/scheduler graph performance related to large-scale shuffle operations.

In order to rearrange data between partitions in dask.dataframe (for parallel merge/sort/shuffle routines), the rearrange_by_column_tasks routine is used to build a task graph for staged shuffling. Since this logic represents n log(n) scaling, the time required for graph creation and execution itself can be quite significant.

Note that a detailed explanation of a nearly identical "staged shuffle" is described in this discussion. One component of the algorithm that is clearly dominating the size of the graph is the repetition of shuffle_group tasks (which output dictionaries of pd/cudf DataFrame objects) and getitem tasks (which select elements of the shuffle-group output). It is my understanding that some people may have promising ideas to improve performance here.

cc @kkraus14 @quasiben @mrocklin (Please do cc others as well..)

[kkraus14]

cc @madsbk who had some ideas in this area

[kkraus14]

Also, taken from #6137 (comment), here's a standalone example of where client/scheduler graph performance is problematic:

from distributed import Client
from dask.datasets import timeseries
from dask.dataframe.shuffle import shuffle


client = Client()

ddf_d = timeseries(start='2000-01-01', end='2005-01-01', partition_freq='1d')  # 1827 partitions
ddf_d_2 = shuffle(ddf_d, "id", shuffle="tasks")

%time ddf_d_2 = ddf_d_2.persist()  #~8s on my machine