Repository: pydata/xarray

Issue:
Dataset combine_by_coords unexpected behavior
### What happened?

I am trying to combine two datasets with partially overlapping data variables and exactly identical coordinates, using the `compat="override"` option.
In the resulting dataset, the data variable, which appears in both datasets contains the values, from the dataset, placed second in the list of datasets to combine.
The result is the same, if I switch the order of the two datasets in the call to `combine_by_coords`

### What did you expect to happen?

I would expect, that for the data variables, which appear in both datasets, the values from the first dataset in the list, passed to `combine_by_coords` are used.

### Minimal Complete Verifiable Example

```Python
import numpy as np
import pandas as pd
import xarray as xr

temperature = np.random.randint(1,255,size=(9,10,10))
precipitation = np.random.randint(1,255,size=(9,10,10))
precipitation_alt = np.random.randint(255,1000,size=(9,10,10))
lon = np.linspace(10,100,10)
lat = np.linspace(10,100,10)
time_0 = pd.date_range("2014-09-06", periods=9, freq='2D')

ds_0 = xr.Dataset(
   data_vars=dict(
       temperature=(('time', 'y', 'x'), temperature),
       precipitation=(('time', 'y', 'x'), precipitation)),
   coords=dict(
       x=lon,
       y=lat,
       time=time_0)
)
ds_1 = xr.Dataset(
    data_vars=dict(
        temperature=(('time', 'y', 'x'), temperature),
        precipitation=(('time', 'y', 'x'), precipitation_alt),
        precipitation2=(('time', 'y', 'x'), precipitation)),
    coords=dict(
        x=lon,
        y=lat,
        time=time_0)

)

res1 = xr.combine_by_coords([ds_0, ds_1], compat="override")
res2 = xr.combine_by_coords([ds_1, ds_0], compat="override")

# In the first case the resulting dataset should contain the same values in the variables precipitation and precipitation2, but it does not.
```


### MVCE confirmation

- [X] Minimal example — the example is as focused as reasonably possible to demonstrate the underlying issue in xarray.
- [X] Complete example — the example is self-contained, including all data and the text of any traceback.
- [X] Verifiable example — the example copy & pastes into an IPython prompt or [Binder notebook](https://mybinder.org/v2/gh/pydata/xarray/main?urlpath=lab/tree/doc/examples/blank_template.ipynb), returning the result.
- [X] New issue — a search of GitHub Issues suggests this is not a duplicate.
- [X] 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?

_No response_

### Environment

<details>
NSTALLED VERSIONS
------------------
commit: None
python: 3.8.18 (default, Sep 11 2023, 13:40:15) 
[GCC 11.2.0]
python-bits: 64
OS: Linux
OS-release: 5.15.0-91-generic
machine: x86_64
processor: x86_64
byteorder: little
LC_ALL: None
LANG: en_US.UTF-8
LOCALE: ('en_US', 'UTF-8')
libhdf5: 1.12.2
libnetcdf: 4.9.3-development

xarray: 2023.1.0
pandas: 2.0.3
numpy: 1.24.4
scipy: 1.10.1
netCDF4: 1.6.5
pydap: None
h5netcdf: None
h5py: None
Nio: None
zarr: None
cftime: 1.6.3
nc_time_axis: None
PseudoNetCDF: None
rasterio: None
cfgrib: None
iris: None
bottleneck: 1.3.5
dask: None
distributed: None
matplotlib: 3.7.2
cartopy: None
seaborn: None
numbagg: None
fsspec: None
cupy: None
pint: 0.21.1
sparse: None
flox: None
numpy_groupies: None
setuptools: 68.2.2
pip: 23.3.1
conda: None
pytest: None
mypy: None
IPython: 8.12.3
sphinx: None
</details>



Discussion:
The root cause of your issue is that the datasets are re-sorted before being combined:

https://github.com/pydata/xarray/blob/cf3655968b8b12cc0ecd28fb324e63fb94d5e7e2/xarray/core/combine.py#L952-L955

The first line re-sorts the data_objects to [ds_1, ds_0] because `vars_as_keys` will give something like this:

`(('precipitation', 'temperature'), ('precipitation', 'precipitation2', 'temperature'))`

and the sorted representation is:

`(('precipitation', 'precipitation2', 'temperature'), ('precipitation', 'temperature'))`

So this will end up in the order [ds_1, ds_0] before entering the underlying combine-processing.

Not sure how this could be resolved. Given the declaration in the docstring this is very surprising behaviour.
It seems like the additional sorting is not needed at all.



Patch:
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index c76f3800f2b..462872bb80e 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -106,6 +106,8 @@ Bug fixes
   By `Kai Mühlbauer <https://github.com/kmuehlbauer>`_.
 - Fix weighted ``polyfit`` for arrays with more than two dimensions (:issue:`9972`, :pull:`9974`).
   By `Mattia Almansi <https://github.com/malmans2>`_.
+- Preserve order of variables in :py:func:`xarray.combine_by_coords` (:issue:`8828`, :pull:`9070`).
+  By `Kai Mühlbauer <https://github.com/kmuehlbauer>`_.
 - Cast ``numpy`` scalars to arrays in :py:meth:`NamedArray.from_arrays` (:issue:`10005`, :pull:`10008`)
   By `Justus Magin <https://github.com/keewis>`_.
 
diff --git a/xarray/core/combine.py b/xarray/core/combine.py
index f2852443d60..f02d046fff6 100644
--- a/xarray/core/combine.py
+++ b/xarray/core/combine.py
@@ -1,8 +1,7 @@
 from __future__ import annotations
 
-import itertools
-from collections import Counter
-from collections.abc import Iterable, Iterator, Sequence
+from collections import Counter, defaultdict
+from collections.abc import Callable, Hashable, Iterable, Iterator, Sequence
 from typing import TYPE_CHECKING, Literal, TypeVar, Union, cast
 
 import pandas as pd
@@ -269,10 +268,7 @@ def _combine_all_along_first_dim(
     combine_attrs: CombineAttrsOptions = "drop",
 ):
     # Group into lines of datasets which must be combined along dim
-    # need to sort by _new_tile_id first for groupby to work
-    # TODO: is the sorted need?
-    combined_ids = dict(sorted(combined_ids.items(), key=_new_tile_id))
-    grouped = itertools.groupby(combined_ids.items(), key=_new_tile_id)
+    grouped = groupby_defaultdict(list(combined_ids.items()), key=_new_tile_id)
 
     # Combine all of these datasets along dim
     new_combined_ids = {}
@@ -606,6 +602,21 @@ def vars_as_keys(ds):
     return tuple(sorted(ds))
 
 
+K = TypeVar("K", bound=Hashable)
+
+
+def groupby_defaultdict(
+    iter: list[T],
+    key: Callable[[T], K],
+) -> Iterator[tuple[K, Iterator[T]]]:
+    """replacement for itertools.groupby"""
+    idx = defaultdict(list)
+    for i, obj in enumerate(iter):
+        idx[key(obj)].append(i)
+    for k, ix in idx.items():
+        yield k, (iter[i] for i in ix)
+
+
 def _combine_single_variable_hypercube(
     datasets,
     fill_value=dtypes.NA,
@@ -965,8 +976,7 @@ def combine_by_coords(
         ]
 
         # Group by data vars
-        sorted_datasets = sorted(data_objects, key=vars_as_keys)
-        grouped_by_vars = itertools.groupby(sorted_datasets, key=vars_as_keys)
+        grouped_by_vars = groupby_defaultdict(data_objects, key=vars_as_keys)
 
         # Perform the multidimensional combine on each group of data variables
         # before merging back together


Test patch:
diff --git a/xarray/tests/test_combine.py b/xarray/tests/test_combine.py
index dfd047e692c..cc20ab414ee 100644
--- a/xarray/tests/test_combine.py
+++ b/xarray/tests/test_combine.py
@@ -1043,6 +1043,20 @@ def test_combine_by_coords_incomplete_hypercube(self):
         with pytest.raises(ValueError):
             combine_by_coords([x1, x2, x3], fill_value=None)
 
+    def test_combine_by_coords_override_order(self) -> None:
+        # regression test for https://github.com/pydata/xarray/issues/8828
+        x1 = Dataset({"a": (("y", "x"), [[1]])}, coords={"y": [0], "x": [0]})
+        x2 = Dataset(
+            {"a": (("y", "x"), [[2]]), "b": (("y", "x"), [[1]])},
+            coords={"y": [0], "x": [0]},
+        )
+        actual = combine_by_coords([x1, x2], compat="override")
+        assert_equal(actual["a"], actual["b"])
+        assert_equal(actual["a"], x1["a"])
+
+        actual = combine_by_coords([x2, x1], compat="override")
+        assert_equal(actual["a"], x2["a"])
+
 
 class TestCombineMixedObjectsbyCoords:
     def test_combine_by_coords_mixed_unnamed_dataarrays(self):