Applying GroupBy.map along chunks instead of dimensions

[sadsimulation]

I'm processing larger than memory datasets using dask-backed xarrays.
However I often need to perform indexing using dataset1.where(array2, drop=True).

The intermediate computations to get dataset1 and array2 are fairly computationally intensive and make use of the xarray indexing and dimension name features, so I feel like xr.apply_ufunc with dask='allowed' wouldn't be a good fit here. Unfortunately the masking causes the output sizes of the output xr.Dataset dimensions to change based on the data (peak detection), so xr.map_blocks is difficult to apply because I don't know the output template shape.

A workaround that I have used to get things working at all is to use dataset.groupby('mydim').map(myfunc) on a dataset that is not backed by dask arrays. This is not great because the groups along mydim vary in size and don't fit a simple uniform chunking along that dimension and predictably only utilizes a single core as the groups are processed sequentially.

Is there an easy way to do something like dataset.chunk(mydim=100).groupby('mydim').map(myfunc) that would utilize my machines CPU better?

[dcherian]

If myfunc handles and returns dask arrays, then the graph construction will happen in serial but the compute will be parallel. Will that work?

Can you create a small example to show how your groups are patterned? It seems like each group is sequential?