diff --git a/bench/duckdb-spike.md b/bench/duckdb-spike.md index f9553089a..8cf6913c8 100644 --- a/bench/duckdb-spike.md +++ b/bench/duckdb-spike.md @@ -31,11 +31,12 @@ uv run python -m bench.run --arms lpspec duckdb \ uv run python -m bench.report bench/results/duckdb-spike.jsonl --arms lpspec duckdb ``` -**Expect the in-tree engine to be slower than the recorded numbers**, and say so -rather than quietly re-baselining: the port has no counterpart to -`labels._factored` (§9), so masked models fall back to the counted path and pay -an ordered window the shipped engine avoided. The committed results file is the -*old* engine and is kept as the measurement the argument was built on. +**The committed results file is the *old* engine**, kept as the measurement the +argument was built on rather than re-baselined. The in-tree engine has since +grown the `_factored` label path §9 recorded as missing — masked models no +longer pay an ordered window over the whole coordinate product — which took the +build on `dispatch/l` from 1.84 s to 0.69 s and on `nodal/l` from 1.09 s to +0.53 s, against 0.59 s and 0.29 s on polars. --- @@ -335,8 +336,11 @@ Where the tax actually falls: column lists on both sides, alias qualification and null-safe `IS NOT DISTINCT FROM` predicates where polars gets all three from column-name semantics. Every fragment composition pays it. -- **`labels._factored`, 37 lines, no counterpart.** The one place polars is - *algorithmically* ahead rather than incidentally, and real work to port. +- **`labels._factored`, 37 lines.** Was the one place polars was + *algorithmically* ahead rather than incidentally. Now ported, and the port + moved the *choice* — which of the three routes a mask allows — up into + `plan.free_prefix`, where it is a question about the plan that both engines + ask once rather than two implementations that could drift. - SQL **wins** on `_compare` (15 → 4) and on the semi-joins, where `WHERE EXISTS` says plainly what `how='semi'` needs a reader to know is not a join. @@ -369,10 +373,10 @@ read-back are written against `ModelTables` and the label frames, so ~2,300-line figure priced a replacement; a *second* engine is a compiler and an assembler. -**Still open, and the reason to re-read §7 before relying on it:** -`labels._factored` has no duckdb counterpart, so masked models fall back to the -counted path and pay an ordered window. The committed results predate the -in-tree port and do not include that cost. +**Read §7 as provenance, not as the current cost.** The committed results +predate the in-tree port, and the port has been measured and improved since — +the factored label path alone halves the build on every masked model. Re-run +the command at the top before quoting a ratio. **Recommended next step, if this is live:** decide the write-path question first. If the answer is no, this closes for good and diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 8ede80846..72b7a7320 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -348,7 +348,9 @@ in the lane is order-free, which is what lets the query planner rearrange it. depending on how much of the product survives the mask — arithmetic, factored, counted — and they must agree integer for integer. That agreement is why labelling is a module with stated inputs rather than three methods among - twenty: nothing else about a build can move an index. + twenty: nothing else about a build can move an index. **Which** of the three + is a property of the plan, not of an engine, so both engines ask + `plan.free_prefix` and only the executions are written twice. - The same order comes **back**: `primal` / `dual` / `to_parquet` sort on the label before handing rows over, the order the LP sink writes. Stated at the read rather than assumed, because a `where` decides which rows survive and a @@ -413,7 +415,7 @@ than discovered at solve time. | `relational/engines/polars/labels.py` | which coordinate gets which solver index; three routes to one number, which must agree | | `relational/engines/polars/executor.py` | assemble the model frames from the bound data | | `relational/result.py` | what a solve returned: status, objective, and the label joins that read values back | -| `relational/sinks/tables.py` | what every sink reads and no more — the four frames plus the batching scalars; what an engine produces | +| `relational/sinks/tables.py` | what every sink reads and no more — the four frames, their dtypes, and the batching scalars; what an engine produces | | `relational/sinks/` | how a built model leaves: `lp_file`, `solver_direct` (one module each, [README](https://github.com/FBumann/lpspec/blob/main/src/lpspec/relational/sinks/README.md)) | | `linopy/__init__.py` | opt-in shim: `build` / `extend` on a `linopy.Model` | | `linopy/loader.py` | data coercion to `xr.Dataset`, master coords | diff --git a/src/lpspec/relational/engines/duck/compiler.py b/src/lpspec/relational/engines/duck/compiler.py index 5ebaed56d..f14ac9e8f 100644 --- a/src/lpspec/relational/engines/duck/compiler.py +++ b/src/lpspec/relational/engines/duck/compiler.py @@ -56,8 +56,14 @@ def lit(value: object) -> str: return 'TRUE' if value else 'FALSE' if value is None: return 'NULL' - if isinstance(value, float) and not math.isfinite(value): - return "'Infinity'::DOUBLE" if value > 0 else "'-Infinity'::DOUBLE" + if isinstance(value, float): + if not math.isfinite(value): + return "'Infinity'::DOUBLE" if value > 0 else "'-Infinity'::DOUBLE" + # `::DOUBLE`, or SQL reads `0.1` as DECIMAL(2,1) — a different number + # from the double the plan holds, and a different one from what the + # polars engine multiplies by. Decimal also propagates: a coefficient + # built from one stays decimal all the way into `obj`. + return f'{value!r}::DOUBLE' return repr(value) @@ -346,7 +352,7 @@ def _variable_fragment(self, name: str) -> TermFragment: # syntax error rather than an empty projection prefix = ''.join(f'{q(d)}, ' for d in dims) table = q(self.variables[name]) - rel = Rel(f'SELECT {prefix}var_label, 1.0 AS coeff FROM {table}', (*dims, 'var_label', 'coeff')) + rel = Rel(f'SELECT {prefix}var_label, 1.0::DOUBLE AS coeff FROM {table}', (*dims, 'var_label', 'coeff')) masked = self.program.variable(name).where is not None presence = Rel(f'SELECT {", ".join(q(d) for d in dims) or q(UNIT)} FROM {table}', dims) if masked else None return TermFragment( diff --git a/src/lpspec/relational/engines/duck/executor.py b/src/lpspec/relational/engines/duck/executor.py index 0eea61118..8692f5bd1 100644 --- a/src/lpspec/relational/engines/duck/executor.py +++ b/src/lpspec/relational/engines/duck/executor.py @@ -33,10 +33,10 @@ if TYPE_CHECKING: from collections.abc import Iterator, Sequence -_COLS = ('col', 'lb', 'ub', 'vtype') -_ROWS = ('row', 'sense', 'rhs') -_MATRIX = ('row', 'col', 'coeff') -_OBJ = ('col', 'coeff') +#: The four frames a sink reads, and their dtypes — both stated by +#: `sinks/tables.py`, which is what reads them. Shared with the polars engine +#: for the reason the frames are: a sink cannot see which engine filled them. +_COLS, _OBJ, _ROWS, _MATRIX = sinks.COLS, sinks.OBJ, sinks.ROWS, sinks.MATRIX class _Labels(Mapping[str, 'pl.LazyFrame']): @@ -91,6 +91,9 @@ def __init__(self) -> None: self._program: plan.Program | None = None self._var_tables: dict[str, str] = {} self._row_tables: dict[str, str] = {} + #: the dims each declared name is read through — what `plan.free_prefix` + #: needs to know which label path a mask allows + self._name_dims: dict[str, tuple[str, ...]] = {} self._var_labels = _Labels(self._con, self._var_tables) self._row_labels = _Labels(self._con, self._row_tables) self._n_cols = 0 @@ -115,10 +118,30 @@ def _view(self, sql: str, prefix: str) -> str: return name def _register(self, name: str, frame: pl.DataFrame) -> str: - # the frame itself, not `to_arrow()`: duckdb reads polars natively, and - # the round-trip through a pyarrow table is what used to drag pandas - # into a runtime that declares it a bridge out and not a dependency - self._con.register(name, frame) + """Copy *frame* into a table of its own, and return the table's name. + + **Copied, not scanned in place.** A registered frame stays a Python + object: duckdb reads it through the buffer protocol from whichever + worker thread the scan lands on, and those threads need the GIL — which + the thread that called `execute` is holding while it waits. On a plan + with several scans of registered frames feeding one pipeline, that + deadlocks outright: `transport/l` reproduces it as a build that sits at + 0% CPU indefinitely, while the same query over copies returns in 0.3 s. + + Copying is not free — the frame and the table are both resident for the + length of one statement — but it buys a build that cannot stall, and + duckdb's own storage for every scan after the first. + + The frame itself, not `to_arrow()`: duckdb reads polars natively, and + the round-trip through a pyarrow table is what used to drag pandas into + a runtime that declares it a bridge out and not a dependency. + """ + source = f'__source {name}__' + self._con.register(source, frame) + try: + self._con.execute(f'CREATE TABLE {q(name)} AS SELECT * FROM {q(source)}') + finally: + self._con.unregister(source) return name def build(self, program: plan.Program, sources: Mapping[str, Any]) -> None: @@ -129,6 +152,7 @@ def build(self, program: plan.Program, sources: Mapping[str, Any]) -> None: is what a caller's data has to survive whoever builds the model. """ self._program = program + self._name_dims = plan.name_dims(program) bound = bind(program, sources) dims = {n: self._register(f'dim_{n}', f.collect()) for n, f in bound.dimensions.items()} params = {n: self._register(f'par_{n}', f.collect()) for n, f in bound.parameters.items()} @@ -170,29 +194,30 @@ def _label_frame( ) -> tuple[str, int]: """The masked coord product of *dims* with a dense *label* from *start*. - Two paths, not the polars side's three. **Unmasked**, a row's label is - its position in the product — arithmetic on the ordinals, no sort and - nothing to count. **Otherwise** it is counted, which costs the ordered - window. The `_factored` middle path is an optimisation rather than a - semantics, and is left out here on purpose: it is the one place the - polars engine is *algorithmically* ahead, so including a half version - of it would flatter this port. + The same three paths the polars engine has, chosen by the same + function. **Unmasked**, a row's label is its position in the product — + arithmetic on the ordinals, no sort and nothing to count. **Masked but + factoring**, the survivors are a rectangle and the label is arithmetic + again over a ranked survivor set (:meth:`_factored`). **Otherwise** it + is counted, which costs the ordered window over the whole product. + + `plan.free_prefix` decides between them for both engines, because a + label is the solver's own column index: two engines choosing routes + independently is how they would come to build different models. """ - rel = self._q.frame(dims, where) - if where is None and not restrictions: - # row-major: the leading dim's ordinal times the width of the rest - terms: list[str] = [] - stride = 1 - for d in reversed(dims): - terms.append(f'{q(_ordinal(d))} * {stride}') - stride *= self._q.cardinality[d] - position = ' + '.join(reversed(terms)) if terms else '0' - cols = ', '.join(q(d) for d in dims) or q(UNIT) - sql = f'SELECT {cols}, ({position} + {start})::BIGINT AS {q(label)} FROM {rel.alias("p")}' - name = self._view(sql, 'lbl') - return name, start + math.prod(self._q.cardinality[d] for d in dims) - - carrier = rel + if not restrictions: + if where is None: + rows = math.prod(self._q.cardinality[d] for d in dims) + cols = ', '.join(q(d) for d in dims) or q(UNIT) + rel = self._q.frame(dims, None) + sql = f'SELECT {cols}, ({self._row_major(dims, start)})::BIGINT AS {q(label)} FROM {rel.alias("p")}' + return self._view(sql, 'lbl'), start + rows + + free = plan.free_prefix(dims, plan.predicate_dims(where, self._name_dims)) + if free: + return self._factored(dims, free, where, label, start) + + carrier = self._q.frame(dims, where) for on, presence in restrictions: keep = ', '.join(f'l.{q(c)}' for c in carrier.columns) on_sql = ' AND '.join(f'l.{q(d)} IS NOT DISTINCT FROM r.{q(d)}' for d in on) @@ -202,6 +227,32 @@ def _label_frame( f'WHERE {on_sql})', carrier.columns, ) + return self._counted(carrier, dims, label, start) + + def _row_major(self, dims: tuple[str, ...], start: int, alias: str = '') -> str: + """Row-major position in *dims*' coordinate product, offset by *start*. + + The trailing dim has stride 1 and every other is the product of the + cardinalities to its right, so the position is a dot product against + the ordinals the frame already carries — no ordering imposed, because + the answer does not depend on the order rows arrive in. The polars twin + is `Labeller.row_major`, and both arithmetic paths reach a label + through it for the reason the label itself is written once. + """ + prefix = f'{alias}.' if alias else '' + terms: list[str] = [] + stride = 1 + for d in reversed(dims): + terms.append(f'{prefix}{q(_ordinal(d))} * {stride}') + stride *= self._q.cardinality[d] + return ' + '.join([*reversed(terms), str(start)]) + + def _counted(self, carrier: Rel, dims: tuple[str, ...], label: str, start: int) -> tuple[str, int]: + """Rank the surviving rows of *carrier*: the ordered window, and a count. + + The general answer and the expensive one — a global sort of the whole + product, which is what the two arithmetic paths exist to avoid. + """ order = ', '.join(q(_ordinal(d)) for d in dims) or '1' cols = ', '.join(q(d) for d in dims) or q(UNIT) sql = ( @@ -209,10 +260,61 @@ def _label_frame( f'FROM {carrier.alias("c")}' ) name = self._view(sql, 'lbl') - got = self._con.execute(f'SELECT count(*) FROM {q(name)}').fetchone() + return name, start + self._height(name) + + def _factored( + self, + dims: tuple[str, ...], + free: int, + where: plan.Predicate, + label: str, + start: int, + ) -> tuple[str, int]: + """Labels for a mask that reads none of the first *free* dims. + + A mask that cannot see the leading dims removes the same coordinates + under every one of their values, so the survivors are a rectangle: the + full product of the leading dims against one surviving set. Ranking + that set costs a window over the *set*, not over the product — on + `dispatch` it ranks 100 generators instead of 10M + ``(snapshot, generator)`` pairs, and the window is what a global sort + made the dominant cost of the build. + + The label is then arithmetic again, through the same + :meth:`_row_major` the unmasked path uses: row-major over the leading + dims, times the width of the surviving set, plus a survivor's rank + within it. That is the same number the window would have counted, + because for each leading coordinate the same survivors appear in the + same order — which is what `tests/test_engine_parity.py` checks by + comparing the built model against the other engine's. + """ + head, kept = dims[:free], dims[free:] + ranked = self._view( + f'SELECT {", ".join(q(d) for d in kept)}, ' + f'(ROW_NUMBER() OVER (ORDER BY {", ".join(q(_ordinal(d)) for d in kept)}) - 1)::BIGINT AS "__rank" ' + f'FROM {self._q.frame(kept, where).alias("k")}', + 'srv', + ) + width = self._height(ranked) + if width == 0: + # nothing survived anywhere, so there is no rectangle to describe. + # The counted path returns the right columns and dtypes for free. + return self._counted(self._q.frame(dims, where), dims, label, start) + + product = self._q.frame(head, None) + picked = ', '.join([*(f'h.{q(d)}' for d in head), *(f's.{q(d)}' for d in kept)]) + position = self._row_major(head, 0, alias='h') + sql = ( + f'SELECT {picked}, (({position}) * {width} + s."__rank" + {start})::BIGINT AS {q(label)} ' + f'FROM {product.alias("h")} CROSS JOIN {q(ranked)} AS s' + ) + rows = math.prod(self._q.cardinality[d] for d in head) * width + return self._view(sql, 'lbl'), start + rows + + def _height(self, table: str) -> int: + got = self._con.execute(f'SELECT count(*) FROM {q(table)}').fetchone() assert got is not None - height = got[0] - return name, start + height + return int(got[0]) # -- declarations --------------------------------------------------- @@ -223,11 +325,11 @@ def _build_variable(self, v: plan.VariableDeclaration) -> pl.DataFrame: labelled = Rel(f'SELECT * FROM {q(name)}', (*v.dims, 'var_label')) bounded = self._q.bounds(labelled, v) - sql = ( - f'SELECT var_label AS col, lb::DOUBLE AS lb, ub::DOUBLE AS ub, ' - f'{lit(v.variable_type)} AS vtype FROM {bounded.alias("b")}' - ) - cols = self._fetch(sql) + sql = f'SELECT var_label AS col, lb::DOUBLE AS lb, ub::DOUBLE AS ub FROM {bounded.alias("b")}' + # `vtype` is attached here rather than selected as a literal in SQL: + # one word per column is one *copy* of that word per row over the wire, + # and the frame's stated dtype is an Enum holding four bytes. + cols = self._fetch(sql).with_columns(pl.lit(v.variable_type, dtype=sinks.VTYPE).alias('vtype')) bad = cols.filter(pl.col('lb').is_null() | pl.col('ub').is_null()).height if bad: raise DataError(null_bounds_message(v.name, bad)) @@ -286,9 +388,10 @@ def _build_constraint(self, c: plan.ConstraintDeclaration) -> tuple[pl.DataFrame return rows, self._fetch(stacked) # `sum` over `(row, col)` is the terminal aggregate — where duplicates # from Sum and GroupSum, which project rather than aggregate, collapse. - return rows, self._fetch( - f'SELECT row, col, sum(coeff) AS coeff FROM ({stacked}) GROUP BY row, col ORDER BY row, col' - ) + # Unordered: every sink sorts the matrix into the order it needs + # (`lp_file` by `(row, col)`, `solver_direct` by `row`), so an ORDER BY + # here is a second sort of the largest frame in the model for nothing. + return rows, self._fetch(f'SELECT row, col, sum(coeff) AS coeff FROM ({stacked}) GROUP BY row, col') def _build_objective(self, o: plan.ObjectiveDeclaration) -> pl.DataFrame | None: """The objective as ``(col, coeff)``, or ``None`` if it has no terms.""" @@ -342,9 +445,7 @@ def close(self) -> None: def _stack(frames: list[pl.DataFrame], columns: tuple[str, ...]) -> pl.DataFrame: kept = [f for f in frames if f.height] if not kept: - return pl.DataFrame( - schema={c: (pl.String if c in ('sense', 'vtype') else pl.Float64) for c in columns} - ).with_columns([pl.col(c).cast(pl.Int64) for c in columns if c in ('row', 'col')]) + return pl.DataFrame(schema={name: sinks.DTYPES[name] for name in columns}) return pl.concat([f.select(columns) for f in kept]) diff --git a/src/lpspec/relational/engines/polars/executor.py b/src/lpspec/relational/engines/polars/executor.py index 579be12ff..81af6960d 100644 --- a/src/lpspec/relational/engines/polars/executor.py +++ b/src/lpspec/relational/engines/polars/executor.py @@ -16,7 +16,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, get_args +from typing import TYPE_CHECKING, Any import polars as pl @@ -37,27 +37,11 @@ from collections.abc import Mapping, Sequence -#: 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') -_OBJ = ('col', 'coeff') -_ROWS = ('row', 'sense', 'rhs') -_MATRIX = ('row', 'col', 'coeff') - -#: The dtype of each of those columns. ``vtype`` is an ``Enum`` over the -#: variable types the plan declares, rather than a string: it holds one word -#: per column and the same handful of words for the whole model, so as a string -#: it stores that word once per row — 0.098 GB of the ``cols`` frame's 0.333 at -#: 9.8M columns, against 0.010 as an Enum. The Enum also makes the vocabulary -#: explicit, so a fourth variable type added to -#: :data:`~lpspec.relational.plan.VariableType` and not reaching here fails -#: where the column is built rather than in whichever sink first compares -#: against a name it does not know. -_DTYPES = { - 'col': pl.Int64, 'row': pl.Int64, - 'lb': pl.Float64, 'ub': pl.Float64, 'rhs': pl.Float64, 'coeff': pl.Float64, - 'sense': pl.String, 'vtype': pl.Enum(get_args(plan.VariableType)), -} # fmt: skip +#: The four frames a sink reads, and their dtypes — both stated by +#: `sinks/tables.py`, which is what reads them. An engine fills the schema; it +#: does not get to have one of its own. +_COLS, _OBJ, _ROWS, _MATRIX = sinks.COLS, sinks.OBJ, sinks.ROWS, sinks.MATRIX +_DTYPES = sinks.DTYPES #: Deprecated. The engine's failures are now split between diff --git a/src/lpspec/relational/engines/polars/labels.py b/src/lpspec/relational/engines/polars/labels.py index 476ab2d10..27fb9c09c 100644 --- a/src/lpspec/relational/engines/polars/labels.py +++ b/src/lpspec/relational/engines/polars/labels.py @@ -13,6 +13,10 @@ arithmetic, factored, counted — and *those* must agree with each other, which is what makes this a module rather than three methods among twenty. +**Which** of the three is a question about the plan, not about polars, so it is +asked of `plan.free_prefix` and answered identically for the duckdb engine. Only +the three executions live here. + Its inputs are stated rather than reached for: a labeller needs the query (to build the masked product), the dimension cardinalities (to do the arithmetic), and the program (to know which dims a mask reads). Nothing else @@ -26,7 +30,6 @@ import polars as pl -from lpspec.errors import LanguageError from lpspec.relational import plan from lpspec.relational.engines.polars.compiler import UNIT, _ordinal @@ -47,13 +50,7 @@ def __init__( ) -> None: self._q = compiler self._card = dimension_cardinality - #: The dims each name in a where is read through — parameters by their - #: ``dims`` and variables by their ``foreach``. A bare name in a where - #: may be either, and the label planner only asks "which dims does this - #: mask touch". One flat mapping, because the language has one flat - #: namespace and the two cannot collide. - self._param_dims: dict[str, tuple[str, ...]] = {p.name: p.dims for p in program.parameters} - self._param_dims.update({v.name: v.dims for v in program.variables}) + self._param_dims = plan.name_dims(program) def frame( self, @@ -93,7 +90,7 @@ def frame( labelled = frame.select(*dims, self.row_major(dims, start).alias(label)) return labelled.collect(engine='streaming'), start + rows - free = _free_prefix(dims, _predicate_dims(where, self._param_dims)) + free = plan.free_prefix(dims, plan.predicate_dims(where, self._param_dims)) if free: return self._factored(dims, free, where, label, start) @@ -191,52 +188,3 @@ def row_major(self, dims: tuple[str, ...], start: int) -> pl.Expr: position = position + pl.col(_ordinal(d)) * stride stride *= self._card[d] return position - - -def _predicate_dims(where: plan.Predicate, param_dims: Mapping[str, tuple[str, ...]]) -> frozenset[str]: - """Which dims *where* reads. - - A parameter is read through its own dims, a dimension comparison through - the dim it names, and a constant reads nothing. Anything unrecognised - raises: there is no such case today, and a new predicate that forgot to - answer here would silently mislabel a model. - """ - if isinstance(where, plan.BooleanConstant): - return frozenset() - if isinstance(where, plan.DimensionComparison): - return frozenset({where.dimension}) - if isinstance(where, (plan.ParameterComparison, plan.ParameterDefined)): - dims = frozenset(param_dims.get(where.parameter, ())) - # a parameter compared against another parameter reads both - value = getattr(where, 'value', None) - if isinstance(value, str) and value in param_dims: - dims |= frozenset(param_dims[value]) - return dims - if isinstance(where, plan.VariableDefined): - # Read through the variable's own foreach, exactly as a parameter is - # read through its dims. `_free_prefix` then keeps its arithmetic path - # for the leading dims this mask cannot see, as for any other predicate. - return frozenset(param_dims.get(where.variable, ())) - if isinstance(where, (plan.And, plan.Or)): - return _predicate_dims(where.left, param_dims) | _predicate_dims(where.right, param_dims) - if isinstance(where, plan.Not): - return _predicate_dims(where.operand, param_dims) - raise LanguageError( - f'{type(where).__name__} is a predicate the label planner does not know how to read; ' - 'add it to _predicate_dims before using it in a where' - ) - - -def _free_prefix(dims: tuple[str, ...], touched: frozenset[str]) -> int: - """How many leading dims the mask does not read. - - Leading, not merely absent: a label follows declaration order, so only a - prefix leaves the surviving set contiguous under each of its coordinates. - Returns 0 when the mask reads the first dim, which is the case that has to - count its survivors the slow way. - """ - free = 0 - while free < len(dims) and dims[free] not in touched: - free += 1 - # every remaining dim is masked-or-not; the split only helps if something is left - return free if free < len(dims) else 0 diff --git a/src/lpspec/relational/plan.py b/src/lpspec/relational/plan.py index d8f8dac86..4ea34bdcd 100644 --- a/src/lpspec/relational/plan.py +++ b/src/lpspec/relational/plan.py @@ -18,9 +18,12 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Literal, TypeVar +from typing import TYPE_CHECKING, Literal, TypeVar -from lpspec.errors import unknown_name_message +from lpspec.errors import LanguageError, unknown_name_message + +if TYPE_CHECKING: + from collections.abc import Mapping ConstraintSense = Literal['==', '<=', '>='] ObjectiveSense = Literal['min', 'max'] @@ -360,6 +363,73 @@ def constraint(self, name: str) -> ConstraintDeclaration: return _declared(self.constraints, name, 'constraint') +def name_dims(program: Program) -> dict[str, tuple[str, ...]]: + """The dims each declared name is read through. + + Parameters by their ``dims`` and variables by their ``foreach``. A bare + name in a ``where`` may be either, and the questions below only ask which + dims it touches. One flat mapping, because the language has one flat + namespace and the two cannot collide. + """ + dims: dict[str, tuple[str, ...]] = {p.name: p.dims for p in program.parameters} + dims.update({v.name: v.dims for v in program.variables}) + return dims + + +def predicate_dims(where: Predicate, dims: Mapping[str, tuple[str, ...]]) -> frozenset[str]: + """Which dims *where* reads, against the mapping :func:`name_dims` builds. + + A parameter is read through its own dims, a dimension comparison through + the dim it names, and a constant reads nothing. Anything unrecognised + raises: there is no such case today, and a new predicate that forgot to + answer here would silently mislabel a model. + """ + if isinstance(where, BooleanConstant): + return frozenset() + if isinstance(where, DimensionComparison): + return frozenset({where.dimension}) + if isinstance(where, (ParameterComparison, ParameterDefined)): + touched = frozenset(dims.get(where.parameter, ())) + # a parameter compared against another parameter reads both + value = getattr(where, 'value', None) + if isinstance(value, str) and value in dims: + touched |= frozenset(dims[value]) + return touched + if isinstance(where, VariableDefined): + # Read through the variable's own foreach, exactly as a parameter is + # read through its dims. `free_prefix` then keeps the arithmetic path + # for the leading dims this mask cannot see, as for any other predicate. + return frozenset(dims.get(where.variable, ())) + if isinstance(where, (And, Or)): + return predicate_dims(where.left, dims) | predicate_dims(where.right, dims) + if isinstance(where, Not): + return predicate_dims(where.operand, dims) + raise LanguageError( + f'{type(where).__name__} is a predicate the label planner does not know how to read; ' + 'add it to predicate_dims before using it in a where' + ) + + +def free_prefix(dims: tuple[str, ...], touched: frozenset[str]) -> int: + """How many leading dims a mask does not read. + + Leading, not merely absent: a label follows declaration order, so only a + prefix leaves the surviving set contiguous under each of its coordinates. + Returns 0 when the mask reads the first dim, which is the case that has to + count its survivors the slow way. + + Static like the rest of this module, and shared for the reason a label is + shared: the engines must not be able to disagree about which coordinate + gets which solver index, and the cheapest way to guarantee that is for + them to ask one function which route to take. + """ + free = 0 + while free < len(dims) and dims[free] not in touched: + free += 1 + # every remaining dim is masked-or-not; the split only helps if something is left + return free if free < len(dims) else 0 + + def divisor_parameters(*expressions: Expression) -> frozenset[str]: """Parameters appearing anywhere in a divisor position. diff --git a/src/lpspec/relational/sinks/__init__.py b/src/lpspec/relational/sinks/__init__.py index 4b446615d..d32e656ca 100644 --- a/src/lpspec/relational/sinks/__init__.py +++ b/src/lpspec/relational/sinks/__init__.py @@ -13,9 +13,15 @@ from lpspec.relational.sinks.highs import build_highs, solve_direct from lpspec.relational.sinks.lp_file import write_lp_file -from lpspec.relational.sinks.tables import ModelTables +from lpspec.relational.sinks.tables import COLS, DTYPES, MATRIX, OBJ, ROWS, VTYPE, ModelTables __all__ = [ + 'COLS', + 'DTYPES', + 'MATRIX', + 'OBJ', + 'ROWS', + 'VTYPE', 'ModelTables', 'build_highs', 'solve_direct', diff --git a/src/lpspec/relational/sinks/tables.py b/src/lpspec/relational/sinks/tables.py index f43281d9c..7261e094d 100644 --- a/src/lpspec/relational/sinks/tables.py +++ b/src/lpspec/relational/sinks/tables.py @@ -7,14 +7,42 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, get_args -from lpspec.relational import chunking +import polars as pl + +from lpspec.relational import chunking, plan if TYPE_CHECKING: from collections.abc import Iterator - import polars as pl + +#: The columns of each frame, in order. +COLS = ('col', 'lb', 'ub', 'vtype') +OBJ = ('col', 'coeff') +ROWS = ('row', 'sense', 'rhs') +MATRIX = ('row', 'col', 'coeff') + +#: And the dtype of each column. Here rather than in an executor because it is +#: what a *sink* reads: two engines filling the same four frames with different +#: types is a difference no sink asked for and none can see coming. +#: +#: ``vtype`` is an ``Enum`` over the variable types the plan declares, rather +#: than a string: it holds one word per column and the same handful of words +#: for the whole model, so as a string it stores that word once per row — +#: 0.098 GB of the ``cols`` frame's 0.333 at 9.8M columns, against 0.010 as an +#: Enum. The Enum also makes the vocabulary explicit, so a fourth variable type +#: added to :data:`~lpspec.relational.plan.VariableType` and not reaching here +#: fails where the column is built rather than in whichever sink first compares +#: against a name it does not know. +DTYPES = { + 'col': pl.Int64, 'row': pl.Int64, + 'lb': pl.Float64, 'ub': pl.Float64, 'rhs': pl.Float64, 'coeff': pl.Float64, + 'sense': pl.String, 'vtype': pl.Enum(get_args(plan.VariableType)), +} # fmt: skip + +#: The variable-type column's dtype, which an engine builds a literal against. +VTYPE = DTYPES['vtype'] @dataclass(frozen=True) diff --git a/tests/test_duck_engine.py b/tests/test_duck_engine.py new file mode 100644 index 000000000..65220727d --- /dev/null +++ b/tests/test_duck_engine.py @@ -0,0 +1,69 @@ +"""What the duckdb engine owes its connection, rather than its caller. + +`test_engine_parity.py` asks whether the two engines answer the same; this asks +about the one thing that is duckdb's alone — that a build's queries read +duckdb's own storage and never reach back into a Python object mid-plan. + +Reached through `LPSPEC_ENGINE`, the switch a caller has, and skipped when the +extra is absent. +""" + +from __future__ import annotations + +import polars as pl +import pytest + +import lpspec as lps +from lpspec.relational import engines + +pytest.importorskip('duckdb') + +MODEL = { + 'dimensions': {'i': {'dtype': 'int', 'values': [0, 1, 2]}}, + 'parameters': {'b': {'dims': ['i']}}, + 'variables': {'x': {'foreach': ['i'], 'bounds': {'lower': 0, 'upper': 'b'}}}, + 'constraints': {'c': {'foreach': ['i'], 'expression': 'x >= b'}}, + 'objectives': {'o': {'sense': 'minimize', 'expression': 'sum(x, over=i)'}}, +} +SOURCES = {'b': pl.DataFrame({'i': [0, 1, 2], 'value': [1.0, 2.0, 3.0]})} + + +@pytest.fixture +def duck(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv(engines.ENV_VAR, 'duckdb') + with lps.build(MODEL, SOURCES) as ex: + assert type(ex).__name__ == 'DuckExecutor', 'the env var did not select the engine' + yield ex + + +def test_every_bound_source_is_a_table_and_not_a_registered_frame(duck) -> None: + """A scan of a Python object needs the GIL the caller is holding. + + `con.register(name, frame)` leaves the frame a Python object that duckdb + reads through the buffer protocol from whichever worker thread the scan + lands on — and those threads need the GIL, which the thread inside + `execute` is holding while it waits for them. A plan with several such + scans in one pipeline deadlocks outright: `transport/l` reproduces it as a + build sitting at 0% CPU indefinitely, while the same query over copies + returns in 0.3 s. + + A registered frame is not in the catalog, so this is checkable without + reproducing the hang — which as a test would be indistinguishable from CI + wedging. + """ + tables = {name for (name,) in duck._con.execute('SELECT table_name FROM duckdb_tables()').fetchall()} + named = {*duck._q.dimensions.values(), *duck._q.parameters.values()} + assert named <= tables, f'{sorted(named - tables)} are registered frames, not tables' + + +def test_nothing_stays_registered_after_a_build(duck) -> None: + """The copy is a step, not a second home for the data. + + A source left registered would keep the caller's frame alive for the life + of the connection — a second copy of every parameter in a process that + chose this engine to hold fewer. + """ + import duckdb + + with pytest.raises(duckdb.CatalogException, match='does not exist'): + duck._con.execute('SELECT * FROM "__source par_b__"') diff --git a/tests/test_engine_parity.py b/tests/test_engine_parity.py index 757575dd5..60a62ac5c 100644 --- a/tests/test_engine_parity.py +++ b/tests/test_engine_parity.py @@ -36,10 +36,19 @@ 'load': pl.DataFrame({'snapshot': [0, 1, 2], 'value': [12.0, 8.0, 20.0]}), } +#: The same model with a mask that *removes* something. `dispatch`'s +#: `where: p_max > 0` reads only `generator`, so both engines label it by +#: arithmetic over a ranked survivor set rather than by counting — and with +#: every `p_max` positive that arithmetic is never asked a question a wrong +#: answer would show up in. Zeroing one generator is what makes the label a +#: claim: `gas` must be column 1 under every snapshot, not column 2. +MASKED = DISPATCH | {'p_max': pl.DataFrame({'generator': ['wind', 'solar', 'gas'], 'value': [10.0, 0.0, 100.0]})} + #: `storage` carries a cyclic `shift` — the operator a second engine is most #: likely to get subtly wrong, and the one the spike calls hardest to port. MODELS = [ ('examples/dispatch.yaml', DISPATCH), + ('examples/dispatch.yaml', MASKED), ('examples/storage.yaml', DISPATCH), ] @@ -95,6 +104,28 @@ def test_both_engines_build_the_same_model(model, sources): ex.close() +@pytest.mark.parametrize(('model', 'sources'), MODELS) +def test_both_engines_produce_the_declared_schema(model, sources): + """Same columns *and* same dtypes — `.equals` above checks only values. + + Two engines can agree on every number and still hand a sink different + types, and a sink reads the frames without asking who filled them. It cost + the duckdb engine a `vtype` of 10M copies of the word `continuous` where + the other holds an `Enum`, and an `obj.coeff` typed `DECIMAL(2,1)` because + SQL reads `1.0` as a decimal — a different number from the double the plan + holds, and one that overflows above 9.9. + """ + from lpspec.relational import sinks + + declared = {'cols': sinks.COLS, 'obj': sinks.OBJ, 'rows': sinks.ROWS, 'matrix': sinks.MATRIX} + for name in ENGINES: + with using(name), lps.build(ROOT / model, sources) as ex: + tables = ex._tables() + for frame, columns in declared.items(): + schema = dict(getattr(tables, frame).schema) + assert schema == {c: sinks.DTYPES[c] for c in columns}, f'{name}: {frame} is not the declared schema' + + @pytest.mark.parametrize(('model', 'sources'), MODELS) def test_both_engines_solve_to_the_same_answer(model, sources): answers = {} diff --git a/tests/test_lp_text.py b/tests/test_lp_text.py index e10656dc2..3bccb3919 100644 --- a/tests/test_lp_text.py +++ b/tests/test_lp_text.py @@ -155,6 +155,18 @@ def test_one_model_writes_the_same_bytes_every_time(tmp_path: Path) -> None: assert len(set(written)) == 1, 'the same model wrote different bytes' + # A second *build*, not a second write: three writes share one matrix, so + # they cannot see an engine that hands its rows back in a different order + # each time. Ordering the matrix inside the engine would hide that; the + # sink carrying its own sort key is what actually makes it true. + for _ in range(2): + with lps.build(schema, data) as rebuilt: + lp = tmp_path / 'rebuilt.lp' + rebuilt.write_lp(lp) + assert hashlib.sha256(lp.read_bytes()).hexdigest() == written[0], ( + 'two builds of one model wrote different bytes' + ) + def test_section_keywords_survive_sections_far_larger_than_a_buffer(tmp_path: Path) -> None: """The sink writes the keywords itself and polars writes the sections.