pydata / xarray

N-D labeled arrays and datasets in Python
https://xarray.dev
Apache License 2.0
3.63k stars 1.09k forks source link

`concat` and `align` flips coordinates order #9777

Open alippai opened 1 week ago

alippai commented 1 week ago

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

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 commented 1 week ago

This is a similar, but fixed issue: https://github.com/pydata/xarray/issues/2811

kmuehlbauer commented 1 week ago

@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 commented 1 week ago

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 commented 1 week ago

Thanks @benbovy for the details.

alippai commented 1 week ago

@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 commented 1 week ago

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 commented 1 week ago

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)