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..8c70c30bd 100644 --- a/src/lpspec/relational/engines/polars/executor.py +++ b/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: 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():