Support DataTree in Xarray's top level functions

[TomNicholas]

Is your feature request related to a problem?

Sometimes you might want to map one of the xarray top-level functions (especially xr.concat or xr.merge) over DataTree objects.

Whilst this could potentially be done manually, we could also imagine generalizing top-level functions to handle this out of the box.

Describe the solution you'd like

For this to work

xr.concat([dt1, dt2], concat_dim='time')

returning a single DataTree, with xr.concat applied to sets of datasets in corresponding nodes.

Describe alternatives you've considered

We could instead not change xarray's top-level functions but still ensure that its relatively easy to achieve using map_over_subtree, i.e.

concat_datatrees = datatree.map_over_subtree(xr.concat)
dt_concatenated = concat_datatrees([dt1, dt2], dim='time')

This would still require generalizing map_over_subtree to understand iterables of DataTree objects though (see zarr-developers/VirtualiZarr#84 (comment)).

Finally we could just not support this at all, in which case the only way for users to concatenate contents of datatrees node-wise is via something like

ds_concatenated = xr.concat([mytree[node].ds for subtree in mytree], dim="time")

but called for every node in the tree.

Additional context

See zarr-developers/VirtualiZarr#84 (comment) for an example of wanting to do this in VirtualiZarr (cc @jonas-spaeth).

This was actually already something we partly discussed in the datatree design meeting (#8747), but I forgot what the conclusion was (do you remember @keewis @flamingbear @owenlittlejohns?).

Checklist

• apply_ufunc- Support DataTree in apply_ufunc #9789
• concat - Support DataTree in xarray.concat #9778
• merge - Support DataTree in xarray.merge #9790
• align - Support DataTree in xarray.align #9791
• broadcast
• combine_by_coords
• combine_nested
• map_blocks

[TomNicholas]

We also have to consider skipping dimensions that are only defined on some levels - see #9778 (comment)

[shoyer]

apply_ufunc would also be a very high-leverage function to which to add DataTree support.

[shoyer]

Let's make small issues for all of these to make these more discoverable?