perf(polars): two group_sums collide only where the coordinates meet, and the matrix leaves in row order

src/lpspec/relational/engines/polars/compiler.py

@@ -104,6 +104,18 @@ class TermFragment:
     a row on the same solver column however they were reshaped, and the
     terminal aggregate that would collapse them has nothing to do.
     """
+    mapping: tuple[tuple[str, ...], ...] = ()
+    """Every operator that moved this fragment's labels, in order.
+
+    A ``GroupSum`` records ``('group', over, coordinate, into)`` and a
+    ``Translate`` ``('shift', dimension, by, wrap)``. Read only to compare two
+    fragments of one variable: identical mappings send a label to the same row,
+    and mappings that differ *only* in a coordinate send it to the same row
+    exactly where those coordinates agree — a question about a dimension table
+    rather than about the model. See :meth:`PolarsCompiler.may_share_a_column`.
+
+    Not a description of the frame, and nothing reads it to build one.
+    """
     presence_dims: tuple[str, ...] | None = None
     """The columns :attr:`presence` is keyed by; ``None`` means :attr:`dims`.
 
@@ -516,7 +528,13 @@ def _sum_fragment(self, p: TermFragment, over: tuple[str, ...], context: str) ->
         # Constructed rather than `replace`d for exactly that: `presence` has to
         # be *dropped* here, and carrying it would be the silent default.
         return TermFragment(
-            keep, frame, p.is_term, p.survives_dropping(dropped), p.label_dims - dropped, variable=p.variable
+            keep,
+            frame,
+            p.is_term,
+            p.survives_dropping(dropped),
+            p.label_dims - dropped,
+            variable=p.variable,
+            mapping=p.mapping,
         )
 
     def _group_fragment(self, p: TermFragment, g: plan.GroupSum, context: str) -> TermFragment:
@@ -544,7 +562,13 @@ def _group_fragment(self, p: TermFragment, g: plan.GroupSum, context: str) -> Te
         # a group is a sum, so §13 applies here as well: absence does not escape
         # it, which is why this constructs rather than `replace`s — see _sum_fragment
         return TermFragment(
-            (*keep, g.into), frame, p.is_term, keyed, _relabel(p.label_dims, g.over, g.into), variable=p.variable
+            (*keep, g.into),
+            frame,
+            p.is_term,
+            keyed,
+            _relabel(p.label_dims, g.over, g.into),
+            variable=p.variable,
+            mapping=(*p.mapping, ('group', g.over, g.coordinate, g.into)),
         )
 
     def _translate_fragment(self, p: TermFragment, s: plan.Translate, context: str) -> TermFragment:
@@ -601,7 +625,13 @@ def remap(source: pl.LazyFrame, carried: list[str]) -> pl.LazyFrame:
             # quietly gone — a different constraint, which is the shape #239
             # removed from masks and #289 removes from shift.
             presence, presence_dims = self._edge(s, card, vacated=False), (s.dimension,)
-        return replace(p, frame=frame, presence=presence, presence_dims=presence_dims)
+        return replace(
+            p,
+            frame=frame,
+            presence=presence,
+            presence_dims=presence_dims,
+            mapping=(*p.mapping, ('shift', s.dimension, str(s.by), str(s.wrap))),
+        )
 
     def _filled_edge(self, s: plan.Translate, card: int, others: list[str], fill: float) -> pl.LazyFrame:
         """``(dims…, cval=fill)`` at every coordinate the shift vacated.
@@ -661,6 +691,52 @@ def _vacated(self, p: TermFragment, s: plan.Translate, card: int, others: list[s
     # assembly helpers used by the executor
     # ------------------------------------------------------------------
 
+    def may_share_a_column(self, a: TermFragment, b: TermFragment) -> bool:
+        """Whether two fragments of one variable can put a row on one column.
+
+        **Distinct variables never do.** Labels are dense and assigned one
+        declaration at a time, so two fragments naming different variables draw
+        from disjoint ranges however either was reshaped (#408). What is left is
+        whether two fragments of *one* variable send some label to the same
+        **row**.
+
+        A label's row is decided by what moved it, so equal
+        :attr:`~TermFragment.mapping` means the same row and a certain
+        collision. Mappings that differ **only in a coordinate** — the network
+        shape, ``group_sum(f, by=to) - group_sum(f, by=from)`` — send it to the
+        same row exactly where those coordinates agree, which is a question
+        about a *dimension table*: is there a line whose ends are one bus? The
+        `line` table is forty rows where the matrix is 12.6M.
+
+        Anything else is answered **yes**. A shift on one side and not the
+        other, a reduction that left them over different dims, a product that
+        broadcast one wider — each changes where a label lands in a way this
+        does not model, and the cost of being wrong is a silently wrong model
+        against the cost of a sort.
+        """
+        if a.variable is None or b.variable is None:
+            return True
+        if a.variable != b.variable:
+            return False
+        if a.dims != b.dims or a.label_dims != b.label_dims:
+            return True
+        if len(a.mapping) != len(b.mapping):
+            return True
+        differing = []
+        for one, other in zip(a.mapping, b.mapping, strict=True):
+            if one == other:
+                continue
+            kind, *rest = one
+            if kind != 'group' or other[0] != 'group' or (rest[0], rest[2]) != (other[1], other[3]):
+                return True
+            differing.append((rest[0], rest[1], other[2]))
+        return all(self._coordinates_meet(over, one, other) for over, one, other in differing)
+
+    def _coordinates_meet(self, dimension: str, one: str, other: str) -> bool:
+        """Whether any label of *dimension* carries the same value in both."""
+        table = self.dimensions[dimension]
+        return bool(table.select((pl.col(one) == pl.col(other)).any()).collect().item())
+
     @staticmethod
     def constant_scalar(p: TermFragment) -> pl.LazyFrame:
         """The const fragment summed per coordinate: ``(dims…, cval)``.

src/lpspec/relational/engines/polars/executor.py

@@ -35,7 +35,9 @@
 from lpspec.relational.result import Result
 
 if TYPE_CHECKING:
-    from collections.abc import Mapping, Sequence
+    from collections.abc import Callable, Mapping, Sequence
+
+    from polars._typing import MaintainOrderJoin
 
 
 #: The four frames a sink reads, as schemas. Stated here because the executor
@@ -113,7 +115,7 @@ def build(self, program: plan.Program, sources: Mapping[str, Any]) -> None:
 
         self._cols = _stack(cols, _COLS)
         self._rows = _stack([r for r, _ in built], _ROWS)
-        self._matrix = _stack([m for _, m in built if m is not None], _MATRIX)
+        self._matrix = _in_row_order(_stack([m for _, m in built if m is not None], _MATRIX))
         self._obj = _stack([objective] if objective is not None else [], _OBJ)
 
     @property
@@ -185,6 +187,18 @@ def _build_constraint(self, c: plan.ConstraintDeclaration) -> tuple[pl.DataFrame
         The terminal aggregate is where duplicates from ``Sum`` and
         ``GroupSum`` — which project rather than aggregate — collapse, and it
         is skipped where nothing can (:func:`_needs_aggregate`).
+
+        **Either way the share leaves ordered by ``row``.** Every sink reads
+        the matrix a row range at a time, so one that is handed an unordered
+        matrix orders it — a second pass over a finished frame, and a second
+        copy of it while the solver's own model is resident. Ordering it here
+        happens inside the pipeline that is already materialising the rows.
+
+        A lone term gets the order from its *join*: the row frame is in label
+        order and ``maintain_order='left'`` keeps it, for less than sorting the
+        result costs. Several terms are several ordered runs stacked, and runs
+        are not an order, so those sort — which is why only the lone term asks
+        the join for anything, a stack having paid for the order regardless.
         """
 
         lhs = self._q.expression(c.lhs, f"constraint '{c.name}' lhs")
@@ -228,8 +242,13 @@ def _build_constraint(self, c: plan.ConstraintDeclaration) -> tuple[pl.DataFrame
             return rows, None
 
         pieces = []
+        carried_order: MaintainOrderJoin | None = 'left' if len(terms) == 1 else None
         for p, sign in terms:
-            placed = frame.join(p.frame, on=list(p.dims), how='inner') if p.dims else frame.join(p.frame, how='cross')
+            placed = (
+                frame.join(p.frame, on=list(p.dims), how='inner', maintain_order=carried_order)
+                if p.dims
+                else frame.join(p.frame, how='cross')
+            )
             pieces.append(
                 placed.select(
                     'row',
@@ -238,8 +257,9 @@ def _build_constraint(self, c: plan.ConstraintDeclaration) -> tuple[pl.DataFrame
                 )
             )
         stacked = pl.concat(pieces)
-        if not _needs_aggregate([fragment for fragment, _ in terms]):
-            matrix = stacked.collect(engine='streaming')
+        if not _needs_aggregate([fragment for fragment, _ in terms], self._q.may_share_a_column):
+            ordered = stacked if len(pieces) == 1 else stacked.sort('row')
+            matrix = ordered.collect(engine='streaming')
             self._check_no_undefined_divisor(f"constraint '{c.name}'", matrix, c.lhs, c.rhs)
             return rows, matrix
 
@@ -282,7 +302,7 @@ def _build_objective(self, o: plan.ObjectiveDeclaration) -> pl.DataFrame | None:
             return None
         pieces = [p.frame.select(pl.col('var_label').alias('col'), pl.col('coeff')) for p in comp.terms]
         stacked = pl.concat(pieces)
-        if _needs_aggregate(comp.terms, projected=True):
+        if _needs_aggregate(comp.terms, self._q.may_share_a_column, projected=True):
             stacked = stacked.group_by('col').agg(pl.col('coeff').sum())
         return stacked.collect(engine='streaming')
 
@@ -489,7 +509,12 @@ def _spanning(solver: str, quantity: str, values: pl.Series | None, expected: in
         )
 
 
-def _needs_aggregate(terms: Sequence[TermFragment], *, projected: bool = False) -> bool:
+def _needs_aggregate(
+    terms: Sequence[TermFragment],
+    may_share: Callable[[TermFragment, TermFragment], bool],
+    *,
+    projected: bool = False,
+) -> bool:
     """Whether stacking *terms* can put two rows on one solver column.
 
     Named for the answer, not the condition: an inverted test here is a wrong
@@ -498,10 +523,14 @@ def _needs_aggregate(terms: Sequence[TermFragment], *, projected: bool = False)
     Two things can put a label twice into the stack, and they are asked
     separately. A single fragment that is not
     :attr:`~lpspec.relational.engines.polars.compiler.TermFragment.keyed`
-    already holds one twice on its own. Two fragments collide only if they
-    carry the *same* variable — ``x + 2 * x`` is one row each and one column —
-    because labels are dense and assigned a declaration at a time, so distinct
-    variables cannot reach a shared one however either fragment was reshaped.
+    already holds one twice on its own. Whether a *pair* can is *may_share*,
+    which answers no for distinct variables — ``x + 2 * x`` is one column and
+    ``x + y`` is two, because labels are dense and assigned a declaration at a
+    time — and then asks whether two fragments of one variable send some label
+    to one **row**. For the network shape,
+    ``group_sum(f, by=to) - group_sum(f, by=from)``, that happens only where a
+    line's two ends are one bus. See
+    :meth:`~lpspec.relational.engines.polars.compiler.PolarsCompiler.may_share_a_column`.
 
     That second half is what makes the ordinary multi-term constraint free.
     ``reserve_up + reserve_down <= p_max`` stacks two fragments, and reading
@@ -530,8 +559,26 @@ def _needs_aggregate(terms: Sequence[TermFragment], *, projected: bool = False)
     """
     if any(not t.survives_dropping(set(t.dims) if projected else set()) for t in terms):
         return True
-    carried = [t.variable for t in terms]
-    return len(set(carried)) != len(carried)
+    return any(may_share(a, b) for i, a in enumerate(terms) for b in terms[i + 1 :])
+
+
+def _in_row_order(matrix: pl.DataFrame) -> pl.DataFrame:
+    """*matrix* with ``row`` known to ascend — checked, not assumed.
+
+    Every constraint orders its own share and they are stacked in declaration
+    order, which is ascending row ranges, so the concatenation is ordered by
+    construction. polars cannot see that through a ``concat``, and a sink that
+    finds the flag missing orders the whole matrix again.
+
+    So the claim is verified and then stated. ``is_sorted`` is a linear scan
+    over a column the frame already holds; the sort behind it is the
+    correctness floor and is expected never to run.
+    """
+    if not matrix.height:
+        return matrix
+    if not matrix['row'].is_sorted():
+        return matrix.sort('row')
+    return matrix.with_columns(pl.col('row').set_sorted())
 
 
 def _has_repeated_entry(matrix: pl.DataFrame) -> bool:

tests/test_compiler.py

@@ -32,6 +32,8 @@
 
 from __future__ import annotations
 
+from dataclasses import replace
+
 import polars as pl
 import pytest
 
@@ -354,6 +356,32 @@ def test_a_variable_and_its_shape_operators_stay_keyed():
         assert compiler().expression(node, 'test').terms[0].keyed, node
 
 
+def test_two_fragments_of_one_variable_are_compared_by_what_moved_their_labels():
+    """A shared variable is necessary for a shared column, and not sufficient.
+
+    Two `group_sum`s of one variable through *different* coordinates of one dim
+    reach the same row only where those coordinates agree, which is a question
+    about the dimension table. Everything else answers yes: the cost of being
+    wrong is a model whose sinks disagree, against the cost of a sort.
+    """
+    q = compiler()
+    by_bus = plan.GroupSum(plan.Variable('p'), over='generator', coordinate='bus', into='bus')
+    one = q.expression(by_bus, 'test').terms[0]
+    same = q.expression(by_bus, 'test').terms[0]
+    assert one.mapping == (('group', 'generator', 'bus', 'bus'),)
+    assert q.may_share_a_column(one, same), 'one mapping, so one row — certain'
+
+    plain = q.expression(plan.Variable('p'), 'test').terms[0]
+    assert q.may_share_a_column(one, plain), 'different dims are not modelled'
+
+    shifted = q.expression(plan.Translate(by_bus, 'snapshot', by=1), 'test').terms[0]
+    assert shifted.mapping[-1][0] == 'shift'
+    assert q.may_share_a_column(one, shifted), 'a shift on one side is not modelled'
+
+    other = q.expression(plan.Variable('p'), 'test').terms[0]
+    assert not q.may_share_a_column(one, replace(other, variable='q')), 'distinct variables never share'
+
+
 def test_a_term_names_its_variable_through_every_operator():
     """What lets two fragments be compared without reading either.
 

tests/test_relational.py

@@ -620,6 +620,62 @@ def test_the_matrix_aggregate_runs_on_what_repeats_and_not_on_what_might():
         assert matrix['coeff'].to_list() == [4.0, 4.0]
 
 
+def _network(self_loop: bool) -> tuple[dict, dict]:
+    """A balance of flows in minus flows out, with or without a line to itself."""
+    model = {
+        'dimensions': {
+            'snapshot': {'dtype': 'int', 'values': [0, 1]},
+            'bus': {'values': ['b0', 'b1']},
+            'line': {'coords': {'from': 'bus', 'to': 'bus'}},
+        },
+        'parameters': {'cap': {'dims': ['line']}, 'load': {'dims': ['snapshot', 'bus']}},
+        'variables': {'f': {'foreach': ['snapshot', 'line'], 'bounds': {'lower': 0, 'upper': 'cap'}}},
+        'constraints': {
+            'balance': {
+                'foreach': ['snapshot', 'bus'],
+                'expression': 'group_sum(f, over=line, by=to) - group_sum(f, over=line, by=from) == load',
+            }
+        },
+        'objectives': {'o': {'sense': 'minimize', 'expression': 'sum(sum(f, over=line), over=snapshot)'}},
+    }
+    ends = ('b0', 'b1') if not self_loop else ('b0', 'b0')
+    sources = {
+        'line': pl.DataFrame({'line': ['l0', 'l1'], 'from': ['b0', ends[0]], 'to': ['b1', ends[1]]}),
+        'cap': pl.DataFrame({'line': ['l0', 'l1'], 'value': [10.0, 10.0]}),
+        'load': pl.DataFrame(
+            {'snapshot': [0, 0, 1, 1], 'bus': ['b0', 'b1', 'b0', 'b1'], 'value': [0.0, 0.0, 0.0, 0.0]}
+        ),
+    }
+    return model, sources
+
+
+def test_two_group_sums_of_one_variable_collide_only_where_the_coordinates_meet():
+    """`by=to` and `by=from` reach one cell exactly on a line to itself.
+
+    Both fragments carry `f`, so counting variables says the aggregate is
+    reachable and every nonzero in the model gets sorted to find out. Which
+    labels they share is decided by the *line* table — two rows here, forty at
+    the `l` rung, against 12.6M nonzeros — so it is asked there.
+
+    The self-loop is the case that must not be optimised away: `l1` leaves and
+    arrives at `b0`, so `+f - f` lands twice on one cell and has to collapse.
+    A matrix holding the same `(row, col)` twice is not a slower model, it is
+    one the sinks disagree about — an LP reader sums the pair and a solver
+    handed duplicate entries is entitled to do either.
+    """
+    for self_loop, expected in ((False, False), (True, True)):
+        model, sources = _network(self_loop)
+        with lps.build(model, sources) as ex:
+            program = lower_program(MathSchema(**model))
+            terms = ex._q.expression(program.constraints[0].lhs, 'test').terms
+            assert len(terms) == 2 and {t.variable for t in terms} == {'f'}
+            assert _needs_aggregate(terms, ex._q.may_share_a_column) is expected, f'self_loop={self_loop}'
+
+            matrix = ex._tables().matrix
+            cells = matrix.select('row', 'col')
+            assert cells.height == cells.unique().height, 'a cell reached the sinks twice'
+
+
 def test_stacking_distinct_variables_asks_for_no_aggregate():
     """The static answer itself, which no end-to-end assertion can reach.
 
@@ -648,10 +704,10 @@ def fragments(expression):
         )
         with PolarsExecutor() as ex:
             ex.build(program, sources)
-            return ex._q.expression(program.constraints[0].lhs, 'test').terms
+            return ex._q.expression(program.constraints[0].lhs, 'test').terms, ex._q.may_share_a_column
 
-    assert not _needs_aggregate(fragments('x + y'))
-    assert _needs_aggregate(fragments('x + 3 * x'))
+    assert not _needs_aggregate(*fragments('x + y'))
+    assert _needs_aggregate(*fragments('x + 3 * x'))
 
 
 def test_the_objective_skips_the_aggregate_only_when_a_column_cannot_repeat():