fix(groupby): align groupers by label, raise on mismatch (#827)

doc/release_notes.rst

@@ -39,6 +39,7 @@ Upcoming Version
 * Freezing an empty constraint group (e.g. an empty ``isel`` slice) no longer raises ``ValueError: cannot reshape array of size 0``. ``Model(freeze_constraints=True)`` and ``Constraint.freeze()`` now round-trip zero-row constraints losslessly.
 * ``Variable.where`` no longer raises ``ValueError: exact match required for all data variable names`` once a solution is attached (after ``Model.solve``) or the variable is fixed. The fill value now covers auxiliary data variables (``solution``, stashed bounds) instead of only ``labels``/``lower``/``upper``.
 * ``LinearExpression.groupby(...).sum()`` with a multi-dimensional ``DataArray`` grouper now reduces over all of the grouper's dimensions on the default (fast) path, instead of leaking one of them into the result.
+* ``LinearExpression.groupby(...).sum()`` now raises when a grouper's labels are reordered or a different set relative to the expression, instead of silently regrouping by position. Reorder the grouper to match the expression's coordinates before grouping. (https://github.com/PyPSA/linopy/issues/827)
 * ``linopy.testing.assert_linequal`` now aligns dimension order before comparing, so mathematically identical expressions built in different orders (e.g. ``x + y`` versus ``y + x``, which inherit different dimension orders from xarray broadcasting) are correctly treated as equal. Genuinely different expressions still fail.
 
 Version 0.8.0

linopy/expressions.py

@@ -215,19 +215,55 @@ def _unstack_multikey(ds: Dataset, dim: str) -> Dataset:
     return ds.unstack(dim, fill_value=LinearExpression._fill_value)
 
 
+def _check_grouper_alignment(group: Any, data: Dataset) -> None:
+    """
+    Ensure an indexed grouper's labels match the data along each shared dim.
+
+    The fast path matches a ``pd.Series``, ``pd.DataFrame`` or ``DataArray``
+    grouper to the expression by position, so a grouper whose coordinates are
+    reordered -- or a different set entirely -- relative to the expression would
+    silently regroup. linopy does not reindex the grouper: it checks that the
+    labels match and raises otherwise, leaving the caller to align the grouper
+    explicitly. A grouper without an index along a dimension has nothing to align
+    by and keeps the positional match.
+    """
+    shared: list[tuple[Hashable, pd.Index]]
+    if isinstance(group, (pd.Series, pd.DataFrame)):
+        shared = [(group.index.name, group.index)]
+    elif isinstance(group, DataArray):
+        shared = [
+            (dim, group.get_index(dim)) for dim in group.dims if dim in group.indexes
+        ]
+    else:
+        return
+    for dim, index in shared:
+        if dim not in data.indexes or index.equals(data.indexes[dim]):
+            continue
+        detail = (
+            "the same labels in a different order"
+            if set(index) == set(data.indexes[dim])
+            else "a different set of labels"
+        )
+        raise ValueError(
+            f"the grouper's labels along dimension {dim!r} do not match the "
+            f"expression's coordinates ({detail}). linopy matches groupers by "
+            f"position and does not reindex; reorder the grouper to the "
+            f"expression's {dim!r} coordinates before grouping."
+        )
+
+
 def _encode_multikey_group(
-    frame: pd.DataFrame, index: pd.Index
+    frame: pd.DataFrame,
 ) -> tuple[pd.Series, tuple[dict, pd.Index]]:
     """
     Encode a multi-key group frame as a single integer-coded Series.
 
     ``_grouped_sum`` groups by one key, so each row's tuple of key values is
-    mapped to an integer. The returned ``(int_map, columns)`` lets
-    :func:`_restore_multikey_index` rebuild the MultiIndex on the result.
+    mapped to an integer. The frame is already aligned to the data (see
+    :func:`_check_grouper_alignment`), so no reindexing is needed. The returned
+    ``(int_map, columns)`` lets :func:`_restore_multikey_index` rebuild the
+    MultiIndex on the result.
     """
-    index_name = frame.index.name
-    frame = frame.reindex(index)
-    frame.index.name = index_name
     int_map = get_index_map(*frame.values.T)
     coded = frame.apply(tuple, axis=1).map(int_map)
     return coded, (int_map, frame.columns)
@@ -352,6 +388,7 @@ def sum(
             )
 
         group = _resolve_group(self.group, self.data)
+        _check_grouper_alignment(group, self.data)
 
         multikey_frame = (
             None if use_fallback else _multikey_value_frame(group, self.data)
@@ -381,9 +418,7 @@ def sum(
 
             multikey_decode = None
             if isinstance(group, pd.DataFrame):
-                group, multikey_decode = _encode_multikey_group(
-                    group, self.data.indexes[group.index.name]
-                )
+                group, multikey_decode = _encode_multikey_group(group)
 
             assert isinstance(group, pd.Series)
             ds = self._grouped_sum(group, data)

test/test_linear_expression.py

@@ -1695,6 +1695,75 @@ def test_linear_expression_groupby_multidim_preserves_extra_dim() -> None:
     assert_linequal(grouped, expr.groupby(groups).sum(use_fallback=True))
 
 
+class TestGroupbyGrouperAlignment:
+    """
+    A ``pd.Series``/``DataArray`` grouper whose labels are reordered or a
+    different set relative to the expression must raise, not silently regroup
+    by position. See https://github.com/PyPSA/linopy/issues/827.
+    """
+
+    @staticmethod
+    def _expr() -> LinearExpression:
+        m = Model()
+        v = m.add_variables(coords=[[0, 1, 2, 3]], dims=["i"], name="v")
+        return 1 * v
+
+    @pytest.mark.parametrize("use_fallback", [True, False])
+    def test_reordered_series_raises(self, use_fallback: bool) -> None:
+        expr = self._expr()
+        s = pd.Series([1, 1, 2, 2], index=pd.Index([0, 1, 2, 3], name="i"), name="g")
+        with pytest.raises(ValueError, match="different order"):
+            expr.groupby(s.iloc[::-1]).sum(use_fallback=use_fallback)
+
+    @pytest.mark.parametrize("use_fallback", [True, False])
+    def test_reordered_dataarray_raises(self, use_fallback: bool) -> None:
+        expr = self._expr()
+        da = xr.DataArray([1, 1, 2, 2], coords={"i": [0, 1, 2, 3]}, name="g")
+        with pytest.raises(ValueError, match="different order"):
+            expr.groupby(da.isel(i=slice(None, None, -1))).sum(
+                use_fallback=use_fallback
+            )
+
+    @pytest.mark.parametrize("use_fallback", [True, False])
+    def test_mismatched_label_set_raises(self, use_fallback: bool) -> None:
+        expr = self._expr()
+        s = pd.Series([1, 1, 2, 2], index=pd.Index([0, 1, 2, 9], name="i"), name="g")
+        with pytest.raises(ValueError, match="different set of labels"):
+            expr.groupby(s).sum(use_fallback=use_fallback)
+
+    def test_reordered_dataframe_raises(self) -> None:
+        # the DataFrame grouper path aligns by label like Series/DataArray
+        expr = self._expr()
+        df = pd.DataFrame({"g": [1, 1, 2, 2]}, index=pd.Index([0, 1, 2, 3], name="i"))
+        with pytest.raises(ValueError, match="different order"):
+            expr.groupby(df.iloc[::-1]).sum()
+
+    def test_reordered_ndim_dataarray_raises(self) -> None:
+        m = Model()
+        v = m.add_variables(coords=[[0, 1], [0, 1]], dims=["i", "j"], name="v")
+        groups = xr.DataArray(
+            [[1, 1], [2, 2]], dims=["i", "j"], coords={"i": [1, 0], "j": [0, 1]}
+        )
+        with pytest.raises(ValueError, match="dimension 'i'"):
+            (1 * v).groupby(groups).sum()
+
+    @pytest.mark.parametrize("use_fallback", [True, False])
+    def test_aligned_grouper_unaffected(self, use_fallback: bool) -> None:
+        # an aligned grouper still groups by label as before
+        expr = self._expr()
+        s = pd.Series([1, 1, 2, 2], index=pd.Index([0, 1, 2, 3], name="i"), name="g")
+        grouped = expr.groupby(s).sum(use_fallback=use_fallback)
+        assert (grouped.data.g == [1, 2]).all()
+        assert grouped.nterm == 2
+
+    def test_unindexed_grouper_matches_positionally(self) -> None:
+        # a grouper without an index along the dim has nothing to align by
+        expr = self._expr()
+        da = xr.DataArray([1, 1, 2, 2], dims=["i"], name="g")  # no "i" coordinate
+        grouped = expr.groupby(da).sum()
+        assert (grouped.data.g == [1, 2]).all()
+
+
 @pytest.mark.parametrize("use_fallback", [True, False])
 def test_linear_expression_groupby_with_name(v: Variable, use_fallback: bool) -> None:
     expr = 1 * v
@@ -1844,15 +1913,16 @@ def test_linear_expression_groupby_with_dataarray(
 
 
 def test_linear_expression_groupby_with_dataframe_non_aligned(v: Variable) -> None:
+    # a DataFrame grouper aligns by label like Series/DataArray: a reordered
+    # index raises instead of silently regrouping. See issue #827.
     expr = 1 * v
     groups = pd.DataFrame(
         {"a": [1] * 10 + [2] * 10, "b": list(range(4)) * 5}, index=v.indexes["dim_2"]
     )
-    target = expr.groupby(groups).sum()
+    expr.groupby(groups).sum()  # aligned: fine
 
-    groups_non_aligned = groups[::-1]
-    grouped = expr.groupby(groups_non_aligned).sum()
-    assert_linequal(grouped, target)
+    with pytest.raises(ValueError, match="different order"):
+        expr.groupby(groups[::-1]).sum()
 
 
 @pytest.mark.parametrize("use_fallback", [True, False])