diff --git a/src/lpspec/relational/engines/polars/compiler.py b/src/lpspec/relational/engines/polars/compiler.py index 594dfce0d..9056fe3b3 100644 --- a/src/lpspec/relational/engines/polars/compiler.py +++ b/src/lpspec/relational/engines/polars/compiler.py @@ -31,7 +31,7 @@ if TYPE_CHECKING: from collections.abc import Callable, Mapping - from polars._typing import JoinStrategy + from polars._typing import JoinStrategy, MaintainOrderJoin from lpspec.relational.engines.polars.binding import BoundSources @@ -188,8 +188,15 @@ def frame(self, dims: tuple[str, ...], where: plan.Predicate | None) -> pl.LazyF def _coordinate_product(self, dims: tuple[str, ...]) -> pl.LazyFrame: """Cross join of the dim tables: labels and ordinals, nothing else.""" + # **Folded in reverse, then projected back.** polars' streaming engine + # walks a cross join right-major, so folding the dims backwards is what + # makes the product arrive in *declaration* row-major order — which is + # label order, and is what lets `cols` be read positionally instead of + # sorted. `labels._in_label_order` verifies that rather than trusting + # it. The projection restores the declared column order the fold + # reversed; only the row order was ever the point. out: pl.LazyFrame | None = None - for d in dims: + for d in reversed(dims): table = self.dimensions[d].select(pl.col('val').alias(d), pl.col('ord').alias(_ordinal(d))) out = table if out is None else out.join(table, how='cross') if out is None: @@ -197,7 +204,7 @@ def _coordinate_product(self, dims: tuple[str, ...]) -> pl.LazyFrame: # has to be real, since a `where` on a scalar declaration filters # this frame and nothing survives a filter. return pl.LazyFrame({UNIT: [0]}) - return out + return out.select(*(c for d in dims for c in (d, _ordinal(d)))) def parameter_join( self, @@ -207,6 +214,7 @@ def parameter_join( alias: str, subject: str, how: JoinStrategy = 'left', + maintain_order: MaintainOrderJoin | None = None, ) -> pl.LazyFrame: """Join *param* onto *frame*, its value column renamed to *alias*. @@ -219,6 +227,12 @@ def parameter_join( caller has to report rather than a row to drop. A mask that cannot be satisfied without the value asks for ``inner`` — see :func:`_certain_parameters`. + + *maintain_order* is asked for only where the result is ordered by + construction and something downstream reads it that way — the bounds, + which become ``cols``. It is not free (+12 ms on a 10M-row join) and it + is far cheaper than restoring the order afterwards (~110 ms to sort the + same frame), so it is passed deliberately rather than defaulted on. """ declaration = self.program.parameter(param) extra = set(declaration.dims) - set(frame_dims) @@ -226,8 +240,8 @@ def parameter_join( raise LanguageError(f'{subject} has dims {sorted(extra)} outside the foreach dims {list(frame_dims)}') table = self.parameters[param].rename({'value': alias}) if not declaration.dims: - return frame.join(table, how='cross') - return frame.join(table, on=list(declaration.dims), how=how) + return frame.join(table, how='cross', maintain_order=maintain_order) + return frame.join(table, on=list(declaration.dims), how=how, maintain_order=maintain_order) # ------------------------------------------------------------------ # predicates (where masks — row absence) @@ -332,7 +346,14 @@ def walk(e: plan.Expression) -> pl.Expr: alias = f'__bound {e.name}__' if alias not in joined: carrier = self.parameter_join( - carrier, e.name, v.dims, alias, f"bound parameter '{e.name}' of variable '{v.name}'" + carrier, + e.name, + v.dims, + alias, + f"bound parameter '{e.name}' of variable '{v.name}'", + # the label frame arrives in label order and `cols` is + # read positionally, so this join may not shuffle it + maintain_order='left', ) joined.add(alias) return pl.col(alias).cast(pl.Float64) diff --git a/src/lpspec/relational/engines/polars/executor.py b/src/lpspec/relational/engines/polars/executor.py index 30867225d..1c5c09a46 100644 --- a/src/lpspec/relational/engines/polars/executor.py +++ b/src/lpspec/relational/engines/polars/executor.py @@ -39,7 +39,7 @@ #: 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. -_COLS = ('col', 'lb', 'ub', 'vtype') +_COLS = ('lb', 'ub', 'vtype') _OBJ = ('col', 'coeff') _ROWS = ('row', 'sense', 'rhs') _MATRIX = ('row', 'col', 'coeff') @@ -159,8 +159,11 @@ def _build_variable(self, v: plan.VariableDeclaration) -> pl.DataFrame: self._blocks[v.name] = (start, labelled.height) bounded = self._q.bounds(labelled.lazy(), v) + # No `col`: the label frame arrives in label order and the bounds join + # maintains it, so a row's *position* is its solver column index. The + # column would be the frame's own row number — 0.32 GB of it at 40M + # columns, held for as long as the model is. cols = bounded.select( - pl.col('var_label').alias('col'), pl.col('lb').cast(pl.Float64), pl.col('ub').cast(pl.Float64), pl.lit(v.variable_type, dtype=_DTYPES['vtype']).alias('vtype'), diff --git a/src/lpspec/relational/engines/polars/labels.py b/src/lpspec/relational/engines/polars/labels.py index 6c219d1ee..0f5424364 100644 --- a/src/lpspec/relational/engines/polars/labels.py +++ b/src/lpspec/relational/engines/polars/labels.py @@ -96,7 +96,7 @@ def frame( frame = self._q.frame(dims, None) rows = math.prod(self._card[d] for d in dims) labelled = frame.select(*dims, self.row_major(dims, start).alias(label)) - return labelled.collect(engine='streaming'), start + rows + return _in_label_order(labelled.collect(engine='streaming'), label), start + rows free = _free_prefix(dims, _predicate_dims(where, self._param_dims)) if free: @@ -167,16 +167,15 @@ def _factored( return empty, start labelled = ( - self._q.frame(head, None) - .select(*head, self.row_major(head, 0).alias('__position')) - .join(survivors.lazy(), how='cross') + survivors.lazy() + .join(self._q.frame(head, None).select(*head, self.row_major(head, 0).alias('__position')), how='cross') .select( *dims, (pl.lit(start, dtype=pl.Int64) + pl.col('__position') * width + pl.col('__rank')).alias(label), ) .collect(engine='streaming') ) - return labelled, start + labelled.height + return _in_label_order(labelled, label), start + labelled.height def row_major(self, dims: tuple[str, ...], start: int) -> pl.Expr: """Row-major position in *dims*' coordinate product, offset by *start*. @@ -198,6 +197,26 @@ def row_major(self, dims: tuple[str, ...], start: int) -> pl.Expr: return position +def _in_label_order(frame: pl.DataFrame, label: str) -> pl.DataFrame: + """*frame* in label order — checked, not assumed. + + **The order is free but it is not ours.** Both arithmetic paths get it from + the emission order of a cross join, and polars' streaming engine walks one + right-major, which is why the product is folded in reverse + (`compiler._coordinate_product`). That is an implementation detail of a + dependency, and a label is arithmetic on ordinals rather than a position — + so a change there would not corrupt a *label*. It would silently stop the + frame being sorted, and `cols` is read positionally. + + So the claim is verified. `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 frame.height and not frame[label].is_sorted(): + return frame.sort(label) + return frame + + def _predicate_dims(where: plan.Predicate, param_dims: Mapping[str, tuple[str, ...]]) -> frozenset[str]: """Which dims *where* reads. diff --git a/src/lpspec/relational/sinks/tables.py b/src/lpspec/relational/sinks/tables.py index 98fc09337..31dc80bb2 100644 --- a/src/lpspec/relational/sinks/tables.py +++ b/src/lpspec/relational/sinks/tables.py @@ -48,13 +48,19 @@ class ModelTables: """The built model, as a sink sees it. - ``cols`` (col, lb, ub, vtype), ``obj`` (col, coeff), ``rows`` (row, sense, + ``cols`` (lb, ub, vtype), ``obj`` (col, coeff), ``rows`` (row, sense, rhs) and ``matrix`` in COO (row, col, coeff). The scalars are what a sink cannot cheaply recover; the objective constant lives outside the frames because it has no column to attach to. ``col`` and ``row`` are dense ``0..n-1``, so they *are* the solver's own indices and no sink builds a mapping. + + **``cols`` carries no ``col``.** It holds one row per column, in label + order, so a row's position is its index and the column would be the frame's + own row number. Every other frame is sparse in the index and keeps it — + ``obj`` most of all, which looks dense on models where every variable has a + cost and is not (0.71 of ``cols`` on `transport`). """ cols: pl.DataFrame @@ -94,12 +100,11 @@ def dense_columns(self, infinity: float) -> DenseColumns: the vectors come back ready to hand over unedited. ``col`` is dense ``0..n-1``, so it *is* the position a value has to end - up at: lining a frame up with the solver's index is a scatter, and - neither the join that fills the objective's gaps nor the sort that - orders the bounds does anything this does not. That pair had to be - collected whole before the first batch could go over, which cost more - than the model does. A column with no row is left free rather than - holding whatever the allocator returned. + up at — and ``cols`` already arrives in that order, one row per column, + so its three vectors are the frame's own and nothing is scattered. + Only ``cost`` still is, because ``obj`` is genuinely sparse: a variable + in no objective term has no row, and is left free rather than holding + whatever the allocator returned. **Nothing textual crosses into numpy.** A polars ``String`` converts by boxing every value as a Python object, so the test against @@ -109,15 +114,14 @@ def dense_columns(self, infinity: float) -> DenseColumns: import numpy as np count = self.column_count - at = self.cols['col'].to_numpy() - lb = _scattered(count, at, self.cols['lb'].to_numpy(), -infinity) - ub = _scattered(count, at, self.cols['ub'].to_numpy(), infinity) - integral = _scattered( - count, at, self.cols.select(pl.col('vtype') != 'continuous').to_series().to_numpy(), False - ) + # `cols` is already the solver's index, so its three vectors need no + # scatter. `copy=True` rather than in place because they are views of + # the frame now: rewriting an infinity through one would edit the built + # model to suit whichever solver asked last. + lb = np.nan_to_num(self.cols['lb'].to_numpy(), copy=True, neginf=-infinity, posinf=infinity) + ub = np.nan_to_num(self.cols['ub'].to_numpy(), copy=True, neginf=-infinity, posinf=infinity) + integral = self.cols.select(pl.col('vtype') != 'continuous').to_series().to_numpy() cost = _scattered(count, self.obj['col'].to_numpy(), self.obj['coeff'].to_numpy(), 0.0) - np.nan_to_num(lb, copy=False, neginf=-infinity, posinf=infinity) - np.nan_to_num(ub, copy=False, neginf=-infinity, posinf=infinity) return lb, ub, cost, integral def dense_rows(self, infinity: float) -> DenseRows: diff --git a/src/lpspec/relational/sinks/writers/lp_file.py b/src/lpspec/relational/sinks/writers/lp_file.py index f809e2222..0317baca9 100644 --- a/src/lpspec/relational/sinks/writers/lp_file.py +++ b/src/lpspec/relational/sinks/writers/lp_file.py @@ -57,9 +57,11 @@ def write_lp_file(model: ModelTables, path: str | Path) -> None: path = Path(path) objective = model.obj.lazy().sort('col').select(_term(pl.col('coeff'), pl.col('col'))) + # `cols` is positional, so the index is added inside the streamed pipeline + # rather than sorted out of a column the model carried all along. bounds = ( model.cols.lazy() - .sort('col') + .with_row_index('col') .select( pl.concat_str( _bound(pl.col('lb'), '-infinity').alias('lb'), @@ -89,7 +91,7 @@ def write_lp_file(model: ModelTables, path: str | Path) -> None: _sink(bounds, f) for variable_type, keyword in (('binary', 'binary'), ('integer', 'general')): - chosen = model.cols.lazy().filter(pl.col('vtype') == variable_type).sort('col') + chosen = model.cols.lazy().with_row_index('col').filter(pl.col('vtype') == variable_type) if chosen.select(pl.len()).collect().item() == 0: continue f.write(f'\n{keyword}\n'.encode()) diff --git a/tests/test_relational.py b/tests/test_relational.py index 8a2c32547..ecf5c0f29 100644 --- a/tests/test_relational.py +++ b/tests/test_relational.py @@ -1254,3 +1254,44 @@ def test_dense_columns_does_not_edit_the_model_it_projects(): assert all(v == float('inf') for v in tables.cols['ub'].to_list()), 'the frame lost its infinities' assert list(second_ub) == [1e100, 1e100], "the second solver got the first solver's infinity" assert list(first) == [0.0, 0.0] + + +@pytest.mark.parametrize('where', [None, 'cap > 0']) +def test_cols_is_positional_so_a_row_index_is_its_solver_column(where): + """Every row of `cols` sits at its own label, masked or not. + + `cols` carries no `col`: it is one row per column in label order, so a + row's position *is* the solver's index. Both arithmetic label paths get + that order from the emission order of a **cross join**, which is a property + of polars rather than of this package — so it is checked here against the + label frame, which knows the answer independently. + + Two dims on purpose: one dim is a scan and cannot be out of order, so a + single-dim model would pass whatever the product did. Bounds are distinct + per coordinate for the same reason — a frame permuted by a mask, a join, or + a change in someone else's engine still has the right *multiset* of bounds + and the wrong row for every one of them. + """ + caps = [ + {'i': i, 'j': j, 'value': 0.0 if where and (i, j) == (1, 'b') else float(10 * i + ord(j))} + for i in range(4) + for j in ('a', 'b', 'c') + ] + model = { + 'dimensions': {'i': {'dtype': 'int', 'values': list(range(4))}, 'j': {'values': ['a', 'b', 'c']}}, + 'parameters': {'cap': {'dims': ['i', 'j']}}, + 'variables': {'x': {'foreach': ['i', 'j'], 'bounds': {'lower': 0, 'upper': 'cap'}}}, + 'constraints': {'c': {'foreach': ['i', 'j'], 'expression': 'x <= cap'}}, + 'objectives': {'o': {'sense': 'maximize', 'expression': 'sum(sum(x, over=j), over=i)'}}, + } + if where: + model['variables']['x']['where'] = where + + with lps.build(model, {'cap': pl.DataFrame(caps)}) as ex: + tables = ex._tables() + assert 'col' not in tables.cols.columns, 'cols carries an index it does not need' + assert tables.cols.height == tables.column_count + + labels = ex._variables['x'].collect().sort('var_label') + expected = labels.join(pl.DataFrame(caps), on=['i', 'j'], how='left')['value'].to_list() + assert tables.cols['ub'].to_list() == expected, 'a bound is attached to the wrong column'