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:
kmuehlbauer: I've fixed the typing issue and replaced the second occurrence here. For the above shown [minibenchmark](https://github.com/pydata/xarray/pull/9070#issuecomment-2151965535) this a slight performance gain. The asv bench doesn't show a performance chain.

The major win is that it fixes #8828. Ready for another review.
kmuehlbauer: Do you see any blockers here @keewis, @TomNicholas?
kmuehlbauer: I've added this to #10002 and will merge this tomorrow, if there are no complaints.
TomNicholas: Sorry for taking so long to reply to this 😓 

> AFAICT, the sorting before groupby makes sure that all Datasets with same variables are grouped together regardless of the variable order within each Dataset. Update: And the position in the object list.

IIRC the point of the groupby is to group datasets into sets whose variables have _common  concatenation dimensions_. See https://github.com/CWorthy-ocean/roms-tools/issues/223#issuecomment-2611052848 for an example of a use case.

> Removing the sorting will result in more groups but doesn't break the tests. Does that mean we are undertesting?. Or is it just fixed by the subsequent merge?

I _think_ as long as those are just subgroups of the original groups (i.e. those created by the code that does sorting) then yes the subsequent merge should fix it. The reason being that any subset of a group containing variables with common concatenation dimensions should also have those same common concatenation dimensions?
kmuehlbauer: > Sorry for taking so long to reply to this 😓

No worries :-) . So, essentially the proposed code does the same groupby, but keeps the initial order. Did you had a chance to test your notebook against this PR? 
kmuehlbauer: Just for the record, the example from above with the according output. We can see that the defaultdict_groupby generates the same groups and keeps them in their natural order. I think this is along the lines the behaviour @keewis mentioned for `more_itertools.bucket` and `toolz.itertoolz.groupby`. So I'd say we can replace the current implementation with the proposed.


```python
import xarray as xr
import collections
import itertools

def vars_as_keys(ds):
    return tuple(sorted(ds))

def groupby_defaultdict(iter, key=lambda x: x):
    idx = collections.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 groupby_itertools(iter, key=lambda x: x):
    iter = sorted(iter, key=vars_as_keys)
    return itertools.groupby(iter, key=vars_as_keys)

x1 = xr.Dataset({"a": (("y", "x"), [[1]]),
                 "c": (("y", "x"), [[1]]),
                 "b": (("y", "x"), [[1]]),}, 
                coords={"y": [0], "x": [0]})
x2 = xr.Dataset({"d": (("y", "x"), [[1]]),
                 "a": (("y", "x"), [[2]]),},
                coords={"y": [0], "x": [0]})
x3 = xr.Dataset({"a": (("y", "x"), [[3]]),
                 "d": (("y", "x"), [[2]]),},
                coords={"y": [0], "x": [1]})
data_objects = [x2, x1, x3]

for vars, ds_with_same_vars in groupby_defaultdict(data_objects, key=vars_as_keys):
    print(vars, tuple(ds_with_same_vars))
print("---------------------------------------------------------------------")
for vars, ds_with_same_vars in groupby_itertools(data_objects, key=vars_as_keys):
    print(vars, tuple(ds_with_same_vars))
```

```python
('a', 'd') (<xarray.Dataset> Size: 32B
Dimensions:  (y: 1, x: 1)
Coordinates:
  * y        (y) int64 8B 0
  * x        (x) int64 8B 0
Data variables:
    d        (y, x) int64 8B 1
    a        (y, x) int64 8B 2, <xarray.Dataset> Size: 32B
Dimensions:  (y: 1, x: 1)
Coordinates:
  * y        (y) int64 8B 0
  * x        (x) int64 8B 1
Data variables:
    a        (y, x) int64 8B 3
    d        (y, x) int64 8B 2)
('a', 'b', 'c') (<xarray.Dataset> Size: 40B
Dimensions:  (y: 1, x: 1)
Coordinates:
  * y        (y) int64 8B 0
  * x        (x) int64 8B 0
Data variables:
    a        (y, x) int64 8B 1
    c        (y, x) int64 8B 1
    b        (y, x) int64 8B 1,)
---------------------------------------------------------------------
('a', 'b', 'c') (<xarray.Dataset> Size: 40B
Dimensions:  (y: 1, x: 1)
Coordinates:
  * y        (y) int64 8B 0
  * x        (x) int64 8B 0
Data variables:
    a        (y, x) int64 8B 1
    c        (y, x) int64 8B 1
    b        (y, x) int64 8B 1,)
('a', 'd') (<xarray.Dataset> Size: 32B
Dimensions:  (y: 1, x: 1)
Coordinates:
  * y        (y) int64 8B 0
  * x        (x) int64 8B 0
Data variables:
    d        (y, x) int64 8B 1
    a        (y, x) int64 8B 2, <xarray.Dataset> Size: 32B
Dimensions:  (y: 1, x: 1)
Coordinates:
  * y        (y) int64 8B 0
  * x        (x) int64 8B 1
Data variables:
    a        (y, x) int64 8B 3
    d        (y, x) int64 8B 2)

```


Patch:
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index c76f3800..462872bb 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 f2852443..f02d046f 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 dfd047e6..cc20ab41 100644
--- a/xarray/tests/test_combine.py
+++ b/xarray/tests/test_combine.py
@@ -1043,6 +1043,20 @@ class TestCombineDatasetsbyCoords:
         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):