Reason or problem

rechunk_no_shuffle (#1067) only accepts xr.DataArray. If you have a Dataset with multiple dask-backed variables you have to loop manually:

ds = ds.map(lambda v: rechunk_no_shuffle(v, target_mb=128))

Proposal

Accept xr.Dataset in addition to xr.DataArray. Iterate over data variables, rechunk each dask-backed one, and return a new Dataset with the same attributes and coordinates.

Usage:

ds = rechunk_no_shuffle(ds, target_mb=128)
# or
ds = ds.xrs.rechunk_no_shuffle(target_mb=128)

Design:

• rechunk_no_shuffle checks isinstance(agg, xr.Dataset) and maps itself over each variable
• Non-dask variables pass through unchanged
• The Dataset .xrs accessor gets a rechunk_no_shuffle method that delegates to the same function

Drawbacks

None obvious. It's a small extension of existing behavior.

Alternatives

Users can call ds.map(...) themselves but that's boilerplate that the library can handle.