perf(engine): cols is positional, and the order is produced rather than repaired

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,16 +188,23 @@ 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:
             # The empty cross join's unit is one row, not nothing — and the row
             # 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,15 +227,21 @@ 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)
         if extra:
             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)

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'),

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.
 

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:

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())

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'