perf: Add maybe_group_terms_polars() helper in common.py that checks for duplicate (labels, vars) pairs before calling group_terms_polars. Use it in both Constraint.to_polars() and LinearExpression.to_polars() to avoid expensive group_by when terms already reference distinct variables

Author: FBumann

linopy/common.py

@@ -449,6 +449,25 @@ def group_terms_polars(df: pl.DataFrame) -> pl.DataFrame:
     return df
 
 
+def maybe_group_terms_polars(df: pl.DataFrame) -> pl.DataFrame:
+    """
+    Group terms only if there are duplicate (labels, vars) pairs.
+
+    This avoids the expensive group_by operation when terms already
+    reference distinct variables (e.g. ``x - y`` has ``_term=2`` but
+    no duplicates). When skipping, columns are reordered to match the
+    output of ``group_terms_polars``.
+    """
+    varcols = [c for c in df.columns if c.startswith("vars")]
+    keys = [c for c in ["labels"] + varcols if c in df.columns]
+    key_count = df.select(pl.struct(keys).n_unique()).item()
+    if key_count < df.height:
+        return group_terms_polars(df)
+    # Match column order of group_terms (group-by keys, coeffs, rest)
+    rest = [c for c in df.columns if c not in keys and c != "coeffs"]
+    return df.select(keys + ["coeffs"] + rest)
+
+
 def save_join(*dataarrays: DataArray, integer_dtype: bool = False) -> Dataset:
     """
     Join multiple xarray Dataarray's to a Dataset and warn if coordinates are not equal.

linopy/constraints.py

@@ -40,9 +40,9 @@
     generate_indices_for_printout,
     get_dims_with_index_levels,
     get_label_position,
-    group_terms_polars,
     has_optimized_model,
     iterate_slices,
+    maybe_group_terms_polars,
     maybe_replace_signs,
     print_coord,
     print_single_constraint,
@@ -622,7 +622,7 @@ def to_polars(self) -> pl.DataFrame:
 
         long = filter_nulls_polars(long)
         if ds.sizes.get("_term", 1) > 1:
-            long = group_terms_polars(long)
+            long = maybe_group_terms_polars(long)
         check_has_nulls_polars(long, name=f"{self.type} {self.name}")
 
         # Build short DataFrame (labels, rhs, sign) without xarray broadcast.

linopy/expressions.py

@@ -60,6 +60,7 @@
     has_optimized_model,
     is_constant,
     iterate_slices,
+    maybe_group_terms_polars,
     print_coord,
     print_single_expression,
     to_dataframe,
@@ -1463,13 +1464,7 @@ def to_polars(self) -> pl.DataFrame:
 
         df = to_polars(self.data)
         df = filter_nulls_polars(df)
-        if df["vars"].n_unique() < df.height:
-            df = group_terms_polars(df)
-        else:
-            # Match column order of group_terms (group-by keys, coeffs, rest)
-            varcols = [c for c in df.columns if c.startswith("vars")]
-            rest = [c for c in df.columns if c not in varcols and c != "coeffs"]
-            df = df.select(varcols + ["coeffs"] + rest)
+        df = maybe_group_terms_polars(df)
         check_has_nulls_polars(df, name=self.type)
         return df