Preserve order of variables in combine_by_coords

doc/whats-new.rst

@@ -53,6 +53,8 @@ Bug fixes
   By `Stephan Hoyer <https://github.com/shoyer>`_.
 - Fix regression in the interoperability of :py:meth:`DataArray.polyfit` and :py:meth:`xr.polyval` for date-time coordinates. (:pull:`9691`).
   By `Pascal Bourgault <https://github.com/aulemahal>`_.
+- Preserve order of variables in :py:func:`xarray.combine_by_coords` (:issue:`8828`, :pull:`9070`).
+  By `Kai Mühlbauer <https://github.com/kmuehlbauer>`_.
 
 Documentation
 ~~~~~~~~~~~~~

xarray/core/combine.py

@@ -1,8 +1,8 @@
 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, Iterable, Iterator, Sequence
 from typing import TYPE_CHECKING, Literal, TypeVar, Union, cast
 
 import pandas as pd
@@ -606,6 +606,21 @@ def vars_as_keys(ds):
     return tuple(sorted(ds))
 
 
+T = TypeVar("T")
+K = TypeVar("K", bound=Hashable)
+
+
+def groupby_defaultdict(
+    iter: list[T], key: Callable[[T], K] = lambda x: x
+) -> 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 +980,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

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):