issue_comments

@ttung xarray 0.14.0 has a new map_blocks function that does what you want. It currently has some limitations on what the user function can do.

I understand that there might be some challenges with returning xarray objects, but it seems like taking xarray objects should be very straightforward. Anything problematic about that?

This would probably be fine as an opt-in option. I'm a little worried that this would be a confusing model for users -- we don't have any other functions that work like this.

Maybe something similar would work here? Xarray would construct a dummy Xarray chunk, apply the user defined function onto that chunk, and then extrapolate metadata out from there somehow.

Yes, this is another possibility, though with xarray there is quite a bit of metadata to extrapolate! We do something similar already in groupby.

The challenge is that with dask's lazy evaluation, we don't know the structure of the returned objects until after evaluating the wrapped functions. So we can't rebuild xarray objects unless we require redundantly specify all the coordinates and attributes from the return values.

Typically in Dask we run the user defined function on an empty version of the data and hope that it provides an appropriately shaped output. If it fails during this process, we ask the user to provide sufficient information for us to populate metadata. Maybe something similar would work here? Xarray would construct a dummy Xarray chunk, apply the user defined function onto that chunk, and then extrapolate metadata out from there somehow.

I'm likely glossing over several important details, but hopefully the general gist of what I'm trying to convey above is somewhat sensible, even if not doable.

I understand that there might be some challenges with returning xarray objects, but it seems like taking xarray objects should be very straightforward. Anything problematic about that?

In the first version of apply_ufunc, I experimented with applying functions that take and return xarray objects.

The challenge is that with dask's lazy evaluation, we don't know the structure of the returned objects until after evaluating the wrapped functions. So we can't rebuild xarray objects unless we require redundantly specify all the coordinates and attributes from the return values.

The alternative would be to make a parallel but eagerly evaluated version of apply_ufunc, e.g., by calling compute on each chunk and then reassembling the result (e.g., ttps://github.com/pydata/xarray/pull/2616). The downside is that this will load your data into memory, but maybe that's acceptable for your purposes.