concat and align flips coordinates order

[alippai]

What happened?

xr.concat() and xr.align() behavior is based on the values of coordinates, the result is surprising

What did you expect to happen?

Consistent behavior, stable coordinates.

Minimal Complete Verifiable Example

This gives the expected (date, c, d):

import pandas as pd
import xarray as xr

ds1 = pd.DataFrame({"date": [1, 1], "c": [1,2], "d": [3,4], "v":[5,6]}).set_index(['date', 'c', 'd']).to_xarray()
ds2 = pd.DataFrame({"date": [2, 2], "c": [1,2], "d": [3,4], "v":[5,6]}).set_index(['date', 'c', 'd']).to_xarray()
xr.concat([ds1, ds2], dim='date')

But changing one value of c results in (c,d,date) coordinates.

import pandas as pd
import xarray as xr

ds1 = pd.DataFrame({"date": [1, 1], "c": [1,2], "d": [3,4], "v":[5,6]}).set_index(['date', 'c', 'd']).to_xarray()
ds2 = pd.DataFrame({"date": [2, 2], "c": [1,3], "d": [3,4], "v":[5,6]}).set_index(['date', 'c', 'd']).to_xarray()
xr.concat([ds1, ds2], dim='date')

I tried to use xr.align instead, but that's similar.
(date, c, d) is coordinates are created:

import pandas as pd
import xarray as xr

ds1 = pd.DataFrame({"date": [1, 1], "c": [1,2], "d": [3,4], "v":[5,6]}).set_index(['date', 'c', 'd']).to_xarray()
ds2 = pd.DataFrame({"date": [2, 2], "c": [1,2], "d": [3,4], "v":[5,6]}).set_index(['date', 'c', 'd']).to_xarray()
xr.align(ds1, ds2, join='outer', exclude='date')

Changing a coordinate results in coordinate order (c,d,date):

import pandas as pd
import xarray as xr

ds1 = pd.DataFrame({"date": [1, 1], "c": [1,2], "d": [3,4], "v":[5,6]}).set_index(['date', 'c', 'd']).to_xarray()
ds2 = pd.DataFrame({"date": [2, 2], "c": [1,3], "d": [3,4], "v":[5,6]}).set_index(['date', 'c', 'd']).to_xarray()
xr.align(ds1, ds2, join='outer', exclude='date')

MVCE confirmation

• Minimal example — the example is as focused as reasonably possible to demonstrate the underlying issue in xarray.
• Complete example — the example is self-contained, including all data and the text of any traceback.
• Verifiable example — the example copy & pastes into an IPython prompt or Binder notebook, returning the result.
• New issue — a search of GitHub Issues suggests this is not a duplicate.
• Recent environment — the issue occurs with the latest version of xarray and its dependencies.

Relevant log output

No response

Anything else we need to know?

Maybe a good workaround or description in the docs is the solution if fixing this behavior is not straightforward.
Right now I have no idea how to append two disjoint Datasets along an axis. transpose changes DataArray coordinates only.

xr.merge() works, but it's really slow (10x slower than concat or align) for appending un-aligned datasets along an axis.

Environment

python: 3.10.14 | packaged by conda-forge OS: Linux xarray: 2024.9.0 pandas: 2.2.2 numpy: 1.26.4

[alippai]

This is a similar, but fixed issue: #2811

[kmuehlbauer]

@alippai Yes, xr.concat uses align under the hood. The issue is with align that it doesn't seem to have the capability to keep the order.

cc @benbovy Do you have a hunch which parts of the alignment process does this reordering?

[benbovy]

I suspect the re-ordering happens in the example above because the "c" coordinate needs to be re-indexed while the other "date" and "d" coordinates do not need to be re-indexed. When all coordinate labels are the same in both ds1 and ds2, no re-ordering happens since no coordinate has to be re-indexed (internal optimization).

I guess that for align we could keep track of the coordinate order in each original dataset and preserve it in the aligned datasets? For concat we would still need to determine which order should prevail since coordinates may be ordered differently in the input datasets. Although I don't have examples at hand, coordinate re-ordering may also happen at other places in Xarray. Perhaps it would be easier (for now) to consider coordinate ordering as an implementation detail and clearly mention where this is not always guaranteed?

[kmuehlbauer]

Thanks @benbovy for the details.

[alippai]

@benbovy thanks for the great details!
I wonder if part of the concat work is done already?

import pandas as pd
import xarray as xr

ds1 = pd.DataFrame({"date": [1, 1], "c": [1,2], "d": [3,4], "v":[5,6]}).set_index(['date', 'c', 'd']).to_xarray()
ds2 = pd.DataFrame({"date": [2, 2], "c": [1,3], "d": [3,4], "v":[5,6]}).set_index(['date', 'c', 'd']).to_xarray()
print(xr.concat([ds1, ds2],dim='date'))

Gives:

<xarray.Dataset> Size: 152B
Dimensions:  (date: 2, c: 3, d: 2)
Coordinates:
  * c        (c) int64 24B 1 2 3
  * d        (d) int64 16B 3 4
  * date     (date) int64 16B 1 2
Data variables:
    v        (date, c, d) float64 96B 5.0 nan nan 6.0 nan ... nan nan nan 6.0

While:

import pandas as pd
import xarray as xr

ds1 = pd.DataFrame({"date": [1, 1], "c": [1,2], "d": [3,4], "v":[5,6]}).set_index(['date', 'c', 'd']).to_xarray()
ds2 = pd.DataFrame({"date": [2, 2], "c": [1,2], "d": [3,4], "v":[5,6]}).set_index(['date', 'c', 'd']).to_xarray()
print(xr.concat([ds1, ds2],dim='date'))

Prints:

<xarray.Dataset> Size: 152B
Dimensions:  (date: 2, c: 3, d: 2)
Coordinates:
  * c        (c) int64 24B 1 2 3
  * d        (d) int64 16B 3 4
  * date     (date) int64 16B 1 2
Data variables:
    v        (date, c, d) float64 96B 5.0 nan nan 6.0 nan ... nan nan nan 6.0

The v variable keeps the coordinates, the dataset coordinates are the ones which are changed.
Is the dataset coordinate metadata only or if it's important for some optimizations, how do I regenerate them?

[kmuehlbauer]

What should be the natural order if we have ds1 and ds2 with concurring layouts? For your example there is a natural order which we might be able to keep track of (this can be fixed, imho). For contradicting layouts we almost everywhere resort to the layout of the first dataset.

[alippai]

Good question.
I scratch only the surface, but I see 3 possibilities:

1. The first dataset
2. The dim= from the concat followed by the remaining from the first dataset
3. The dimensions, the shape of the result (which was always (date, c, d) above)