Dask computes more than necessary during groupby

user_1785661 0
Viewed 51

It seems to me that Dask does more work than necessary when using DataFrame.groupby.

If I set up some data and want only a subset of the resulting operation,

import dask
df = dask.datasets.timeseries(partition_freq='2W').set_index('id').persist()
task = df.groupby('id')['x'].sum().head(compute=False)

Dask seems to want to use all partitions for the computation:

task.visualize(optimize_graph=True, size='8')

enter image description here

Groupby automatically collapses the dataframe into a single partition and I suppose head cannot change the graph retrospectively in that regards.

I hoped that perhaps split_out could help this, but that's also not the case.

task = df.groupby('id')['x'].sum(split_out=2).head(compute=False)
task.visualize(optimize_graph=True, size='8')

enter image description here

It therefore seems that in this particular case, Dask does (potentially significantly) more work than necessary.

Question

My questions thus are

  • Is my assessment of the situation correct? Does Dask indeed do much more work? If not, how could I tell?
  • If I'm correct, what can I do to tell Dask to only compute what is needed?

My actual problem is of course much more complex, but I'm hoping these toy examples will help illuminate the problems in my case.

dask dataframe python
1 Answers

It seems that map_partitions is your friend even for groupby-with-index operations!

import dask

# Warning: might be too big, scale according to your machine
df = (dask.datasets.timeseries(end='2002-01-30', partition_freq='2W')
      .set_index('id')
      .persist())

task_a = df.groupby('id')['x'].sum().head(compute=False)
%time task_a.compute()
CPU times: user 22.5 ms, sys: 843 µs, total: 23.3 ms
Wall time: 161 ms

task_b = df.map_partitions(lambda x: x.groupby('id')['x'].sum()).head(compute=False)
%time task_b.compute()
CPU times: user 11.6 ms, sys: 0 ns, total: 11.6 ms
Wall time: 19.6 ms

Results in this much more reasonable graph (shown only for the original data with 2 partitions)

enter image description here

user_17856610
Related