Non-aggregating grouped operations on dask arrays are painfully slow to construct

[shoyer]

These are both entirely lazy operations:

>>> %time res = ds.groupby('time.month').mean('time')
CPU times: user 142 ms, sys: 20.3 ms, total: 162 ms
Wall time: 159 ms
>>> %time res = ds.groupby('time.month').apply(lambda x: x - x.mean())
CPU times: user 46.1 s, sys: 4.9 s, total: 51 s
Wall time: 50.4 s

I suspect the issue (in part) is that _interleaved_concat_slow indexes out single elements from each dask array along the grouped axis prior to concatenating them together (unit tests for interleaved_concat can be found here). So we end up creating way too many small dask arrays.

Profiling results on slightly smaller data are in this gist.

It would be great if we could figure out a way to make this faster, because these sort of operations are a really nice show case for xray + dask.

CC @mrocklin in case you have any ideas.

[mrocklin]

Interesting. Looks like the overhead of managing the graph is becoming non-trivial.

How long do these operations take to execute? What is your blockdims?

[shoyer]

Here's the timing info:

%time res = ds.t2m.groupby('time.month').mean('time').sum()
# CPU times: user 133 ms, sys: 6.39 ms, total: 140 ms
# Wall time: 145 ms
%time res.load_data()
# CPU times: user 2min 47s, sys: 1min, total: 3min 48s
# Wall time: 1min 19s

%time res = ds.t2m.groupby('time.month').apply(lambda x: x - x.mean()).sum()
# CPU times: user 49.1 s, sys: 6.39 s, total: 55.5 s
# Wall time: 55.1 s
%time res.load_data()
# CPU times: user 6min 17s, sys: 2min 20s, total: 8min 38s
# Wall time: 3min 25s

Blocks shape is {'latitude': (256,), 'longitude': (512,), 'time': (124, 124, 112, ..., 124, 120, 124)}

[shoyer]

The good news about that timing info is that dask is still much faster for calculating the graph than doing the actual computation. But it's still not ideal from an interactivity perspective.

[mrocklin]

Hrm, not that much faster. It's 25% overhead on the apply. Worth checking out. I might not get to this before release though.

[shoyer]

Yeah, like I said in the other issue I don't think this is a blocker (we can add a disclaimer to the docs).

[stale]

In order to maintain a list of currently relevant issues, we mark issues as stale after a period of inactivity

If this issue remains relevant, please comment here or remove the stale label; otherwise it will be marked as closed automatically

[jhamman]

I reran @shoyer's original benchmark. I think this can be closed. Constructing the graphs for the example shown above (on slightly different data) took 84.7 ms and 241 ms respectively. Much better than before.