Joining on subset of coords

[JulienBrn]

Hi, I'm looking for an efficient way of joining on subsets of coordinates without passing through pandas. I have tried xr.merge, xr.align, xr.combine_by_coords with little success, but perhaps I'm missing something...

Below is an example where I'm using pandas to achieve the desired result. How could I get the same result using xarray directly ?

import xarray as xr, numpy as np, pandas as pd

d1 = xr.DataArray(np.random.rand(5, 100), dims=["x", "y"])
d2 = xr.DataArray(np.random.rand(4), dims=["x"])

d1["c1"] = xr.DataArray([0, 0, 0, 2, 3], dims="x")
d1["c2"] = xr.DataArray([0, 0, 1, 1, 1], dims="x")
d1["c3"] = xr.DataArray([2, 0, 0, 0, 0], dims="x")

d2["c1"] = xr.DataArray([0, 0, 2, 3], dims="x")
d2["c2"] = xr.DataArray([0, 1, 1, 1], dims="x")


d1= d1.to_dataset(name="d1")
d2= d2.to_dataset(name="d2")

expected_result = d1.to_dataframe().reset_index().merge(d2.to_dataframe().reset_index(drop=True), on=["c1", "c2"], how="left")\
  .set_index(["x", "y"]).to_xarray().set_coords(["c1", "c2", "c3"]).drop_vars(["x", "y"])

#Edit: added the following to remove broadcasting
for a in ["c1", "c2", "c3", "d2"]:
    expected_result[a] = expected_result[a].isel(y=0)

Note, I believe I have gone through the combine doc without finding a similar use case.

[dcherian]

Is this right? It's not broadcasted out like your expected_result though

import xarray.indexes

d1i = d1.set_xindex(("c1", "c2"), xr.indexes.PandasMultiIndex)
d2i = d2.set_xindex(("c1", "c2"), xr.indexes.PandasMultiIndex)

xr.merge([
    d1i, 
    # automatic alignment does not work because of the duplicated (c1, c2) = (0,0) in d1.
    # so we reindex manually
    d2i.reindex_like(d1i)
]).drop_indexes(("x", "c1", "c2"))

[JulienBrn]

Thanks, this seems perfect (I just added the drop_vars("x")).

I had hoped for a more straightforward function as I'm slowly using xarray instead of pandas for almost everything: it almost has all the features pandas has, handles multiple dimensions, provides a usually clearer and more concise display, the API is nicer (I love .sel instead of brackets), natural Dask integration, apply_ufunc as a general solution, datatrees, ... and so many other things. Thank you so much for the great work.

However, joining (as in this example, but would this solution generalize to coordinates on multiple dimensions?) and grouping (though much improved now) seems much less straightforward than in pandas, though I do not see a clear reason why...

[dcherian]

Joining is just more explicit and tied to an Index (like here, you wanted to join on ("c1", "c2") so we had to add an index for that).

grouping (though much improved now) seems much less straightforward than in pandas

Do you have some examples of what could be improved?

[JulienBrn]

Joining passing indexes seems very logical, I would expect something like:

• Set as index c1, c2, c3 in d1 and c1, c2 in d2.
• expect a function like xr.join([stuff_to_join], on=[indexes or subindexes of a multiindex], how=[left, right, inner or outer]) that only joins (merge does so many more things).

But here, I need to create the exact same index, reindex manually, and call merge.

On a more broad question, I would argue that indexes should be used for speed, not semantics (i.e all code should be able to work without having to worry about indices and one should care about them only for performance). In this case, the "on" parameter would be coordinates and behind the scenes an index would be created if not already existing. The fact that pandas has a different API for indexes than other columns is part of what makes the pandas API worse than polars API.

For grouping, I'll come with an example in a couple of days.