issue_comments

Fixed on main with ds.groupby("year").mean(method="blockwise")

I think the behaviour in Ryan's most recent comment is a consequence of groupby.mean being python results = [] for group_idx in group_indices: # one group per year group = ds.isel(group_idx) # (SPLIT) results.append(group.mean()) # (APPLY) return xr.concat(results, dim="year") # COMBINE results in one chunk per year (one chunk per element in results)

I think the fundamental question is: Is it really possible for dask to recognize that the chunk structure after the combine step could be consolidated with an arbitrary number of apply steps in the middle ? OR When a computation maps a single chunk to many chunks, should dask consolidate the output chunks (using array.chunk-size)?

We can explicitly ask for consolidation of chunks by saying the output should be chunked 5 along year python dask.config.set({"optimization.fuse.ave-width": 6}) # note > 5 ( ds.foo.groupby("year") .mean(dim="time") .chunk({"year": 5}) # really important, why and how would dask choose this automatically/ .data.visualize(optimize_graph=False) )

Then if we set optimization.fuse.ave-width appropriately, we get the graph we want after optimization python dask.config.set({"optimization.fuse.ave-width": 6}) ( ds.foo.groupby("year") .mean(dim="time") .chunk({"year": 5}) # really important .data.visualize(optimize_graph=True) )

Can we make dask recognize that the 5 getitem tasks from input-chunk-0, at the bottom of each tower, can be fused to a single task? In that case, fuse the 5 getitem tasks and "propagate" that fusion up the tower.

I guess another failure here is that when fuse.ave-width is 3 (< width of tower), why isn't dask fusing to make three "sub-towers" per-tower? Even that would help reduce number of tasks. dask.config.set({"optimization.fuse.ave-width": 3}) ( ds.foo.groupby("year") .mean(dim="time") .chunk({"year": 5}) # really important .data.visualize(optimize_graph=True) )

Reading up on fusion, the docstring says

This optimization applies to all reductions–tasks that have at most one dependent–so it may be viewed as fusing “multiple input, single output” groups of tasks into a single task.

So we need the opposite : fuse "single input, multiple output" to a single task when some appropriate heuristic is satisfied.

Can this be closed or is there something to do on the xarray side now that dask/dask#3648 has been merged?