From 42cfd85338220f40bff0486862629ac9fc83a217 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Fri, 31 Jul 2026 22:27:31 +0200 Subject: [PATCH 1/2] perf(polars): two group_sums of one variable collide only where the coordinates meet MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #408 established that two fragments can share a solver column only if they carry the same variable. Sharing a variable is necessary and not sufficient, and the assembly treated it as both — so `transport`'s balance, `group_sum(f, by=to) - group_sum(f, by=from)`, sorted every nonzero in the model to discover that nothing repeated. At `l` that is 0.20 s of a 0.69 s constraint, and the answer is always no. Two fragments of one variable reach the same row only if the operators that moved their labels send them there. A fragment now records those operators, so mappings that differ **only in a coordinate** ask the question they actually pose: is there a line whose two ends are one bus? That is a scan of the `line` table — forty rows at `l`, against 12.6M nonzeros. Everything else answers 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 moves a label in a way this does not model, and the cost of being wrong is a matrix holding one cell twice, which an LP reader sums and a solver handed duplicates may do either way. `may_share_a_column` also answers *no* for distinct variables, which is #408's half of the same question, so the assembly asks one thing rather than two. **Not a win on its own** — it is measured with the row-ordering that follows, because dropping the sort alone hands the sink an unordered matrix and it sorts instead. Co-Authored-By: Claude Opus 5 (1M context) --- .../relational/engines/polars/compiler.py | 82 ++++++++++++++++++- .../relational/engines/polars/executor.py | 28 ++++--- tests/test_compiler.py | 28 +++++++ tests/test_relational.py | 62 +++++++++++++- 4 files changed, 184 insertions(+), 16 deletions(-) diff --git a/src/lpspec/relational/engines/polars/compiler.py b/src/lpspec/relational/engines/polars/compiler.py index cde6ef8c1..b6e1d01c5 100644 --- a/src/lpspec/relational/engines/polars/compiler.py +++ b/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)``. diff --git a/src/lpspec/relational/engines/polars/executor.py b/src/lpspec/relational/engines/polars/executor.py index b86c6e1b7..664ea8f07 100644 --- a/src/lpspec/relational/engines/polars/executor.py +++ b/src/lpspec/relational/engines/polars/executor.py @@ -35,7 +35,7 @@ from lpspec.relational.result import Result if TYPE_CHECKING: - from collections.abc import Mapping, Sequence + from collections.abc import Callable, Mapping, Sequence #: The four frames a sink reads, as schemas. Stated here because the executor @@ -238,7 +238,7 @@ def _build_constraint(self, c: plan.ConstraintDeclaration) -> tuple[pl.DataFrame ) ) stacked = pl.concat(pieces) - if not _needs_aggregate([fragment for fragment, _ in terms]): + if not _needs_aggregate([fragment for fragment, _ in terms], self._q.may_share_a_column): matrix = stacked.collect(engine='streaming') self._check_no_undefined_divisor(f"constraint '{c.name}'", matrix, c.lhs, c.rhs) return rows, matrix @@ -282,7 +282,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 +489,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 +503,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 +539,7 @@ 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 _has_repeated_entry(matrix: pl.DataFrame) -> bool: diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 443b910ad..63d0c8754 100644 --- a/tests/test_compiler.py +++ b/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. diff --git a/tests/test_relational.py b/tests/test_relational.py index 9e425beaa..aa56d781d 100644 --- a/tests/test_relational.py +++ b/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(): From cd65349cd580541c155ea3b28c06dfb12ca9fe8e Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Fri, 31 Jul 2026 22:30:01 +0200 Subject: [PATCH 2/2] perf(polars): the matrix leaves the build in row order MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every sink reads the matrix a row range at a time, so one 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 in the build 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 since #433, and `maintain_order='left'` keeps it — half the cost of sorting the result on `dispatch/l` (0.068 s against 0.129 s). Several terms are several ordered runs stacked, and runs are not an order, so those sort. Only the lone term asks the join for anything, because a stack pays the sort regardless and the join's guarantee is then dead weight. `_in_row_order` states the claim after the stack, the way `_in_label_order` does for the labels: the concatenation is ordered by construction — ascending row ranges of ordered shares — but polars cannot see that through a `concat`, so it is verified with a linear scan and then flagged. The sort behind it is the correctness floor and is expected never to run. This is what makes the coordinate test in the previous commit pay. Measured together on `transport/l`, in-process and paired: today: sort(row,col) + check build 0.975s handoff 0.488s total 1.463s coordinate test alone build 0.466s handoff 0.956s total 1.422s both build 0.686s handoff 0.531s total 1.217s Co-Authored-By: Claude Opus 5 (1M context) --- .../relational/engines/polars/executor.py | 45 +++++++++++++++++-- 1 file changed, 42 insertions(+), 3 deletions(-) diff --git a/src/lpspec/relational/engines/polars/executor.py b/src/lpspec/relational/engines/polars/executor.py index 664ea8f07..8c70c30bd 100644 --- a/src/lpspec/relational/engines/polars/executor.py +++ b/src/lpspec/relational/engines/polars/executor.py @@ -37,6 +37,8 @@ if TYPE_CHECKING: 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 #: is what fills them and an empty model still has to have them. @@ -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', @@ -239,7 +258,8 @@ def _build_constraint(self, c: plan.ConstraintDeclaration) -> tuple[pl.DataFrame ) stacked = pl.concat(pieces) if not _needs_aggregate([fragment for fragment, _ in terms], self._q.may_share_a_column): - matrix = stacked.collect(engine='streaming') + 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 @@ -542,6 +562,25 @@ def _needs_aggregate( 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: """Whether a matrix sorted by ``(row, col)`` holds one cell twice.