From 432b534bbc46820ad4d7bd076dcf01afa14b94e2 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Sun, 26 Jul 2026 18:45:27 +0200 Subject: [PATCH] refactor: three modules in the engine, one per box in the diagram MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ARCHITECTURE.md's pipeline has always drawn the relational lane as a plan, an executor and two sinks. The code had one 1,019-line file and a `DuckdbExecutor` with 32 methods that owned duckdb, HiGHS and LP serialisation between them. This makes the picture true. relational/compiler.py plan -> SQL text. Pure: no connection. relational/executor.py owns the database; fills the tables. relational/sinks/ drains them: lp_file, solver_direct. The seam was already there to be found. Every method that built SQL — `_compile`, `_pred_sql`, `_bound_sql`, `_frame_sql`, the three fragment rewriters — read exactly three things: the program, each dimension's cardinality, and which parameters are boolean. None of them touched the connection. `SqlCompiler` is those three fields and those methods, so compiling a plan node now needs no engine at all. That is worth more than tidiness. ARCHITECTURE.md asks you to judge a new operator by reading the SQL it emits — pointwise, bounded-halo or global. Until now that meant building a model and solving it. `tests/test_compiler.py` does it in 19 assertions with nothing installed: hand it `Translate` and check the dim table is joined twice and no window appears; hand it `GroupSum` and check the dim tuple changed and no `GROUP BY` did. `test_every_ir_expr_node_is_handled_by_the_executor` becomes `test_every_plan_node_is_handled_by_the_compiler` and greps the module that actually consumes the nodes — the executor no longer mentions most of them, and a node the compiler ignores has no relational meaning however much the executor moves around it. Sinks are a package, one module per sink, because that is where the fences are: `highspy` is an optional dependency of `solver_direct` alone, and a caller who only writes LP files should not import it (verified — importing `sinks.lp_file` leaves `highspy` out of `sys.modules`). Splitting by kind instead — writers together, solvers together — would put two optional imports in one module and a function that branches on which solver you meant. `sinks/README.md` carries the contract, the checklist for adding one, and why `mps` is deliberately not pre-grouped with `lp_file` yet. `test_architecture_doc_mentions_every_module` becomes `test_every_module_is_documented_somewhere`: a README beside the code counts. Otherwise a subpackage with one member per variant pushes its whole membership list into ARCHITECTURE.md's map, which is the thing that map exists not to be. Pure move: no SQL text, no numerics and no error message changed. LP output is identical block-for-block (it is not byte-stable run to run, on main either — #109). executor.py 1019 -> 595 lines. 308 tests pass, 19 of them new. Co-Authored-By: Claude Opus 5 (1M context) --- ARCHITECTURE.md | 21 +- linopy_yaml/relational/compiler.py | 373 +++++++++++++++ linopy_yaml/relational/executor.py | 548 +++-------------------- linopy_yaml/relational/sinks/README.md | 69 +++ linopy_yaml/relational/sinks/__init__.py | 22 + linopy_yaml/relational/sinks/highs.py | 94 ++++ linopy_yaml/relational/sinks/lp_file.py | 102 +++++ linopy_yaml/relational/sinks/tables.py | 48 ++ tests/test_architecture.py | 45 +- tests/test_compiler.py | 185 ++++++++ 10 files changed, 995 insertions(+), 512 deletions(-) create mode 100644 linopy_yaml/relational/compiler.py create mode 100644 linopy_yaml/relational/sinks/README.md create mode 100644 linopy_yaml/relational/sinks/__init__.py create mode 100644 linopy_yaml/relational/sinks/highs.py create mode 100644 linopy_yaml/relational/sinks/lp_file.py create mode 100644 linopy_yaml/relational/sinks/tables.py create mode 100644 tests/test_compiler.py diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 71599e41..71b44832 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -31,10 +31,11 @@ flowchart TB subgraph REL["Relational lane — streaming · memory-bounded · linopy-free"] direction TB LOWER["lowering.py"] --> PLAN["logical plan
(relational/plan.py)"] - DR[("data
parquet paths / Arrow tables")] --> EXEC - PLAN --> EXEC["executor.py
tidy tables in file-backed duckdb
under memory_limit"] - EXEC --> LPS["lp_file sink
portability, debugging
(mps planned)"] - EXEC --> DIRECT["solver_direct sink
COO batches → highspy → HiGHS"] + PLAN --> COMP["compiler.py
plan → SQL text
pure: no connection"] + DR[("data
parquet paths / pandas")] --> EXEC + COMP --> EXEC["executor.py
tidy tables in file-backed duckdb
under memory_limit"] + EXEC --> LPS["lp_file sink (sinks.py)
portability, debugging
(mps planned)"] + EXEC --> DIRECT["solver_direct sink (sinks.py)
COO batches → highspy → HiGHS"] DIRECT --> SOL["solution tables
(label join, never dense)"] end @@ -195,6 +196,14 @@ sink, is [ROADMAP Track 4](ROADMAP.md#track-4--sink-capabilities). ## The relational lane +**Three modules, one per box above.** `compiler.py` turns plan nodes into SQL +and holds no connection; `executor.py` owns the database and fills the tables; +`sinks.py` drains them. The split is what makes the admissibility test below +something you can perform rather than reason about — build a `SqlCompiler`, +hand it a node, read the `SELECT` (`tests/test_compiler.py` does exactly that, +with no engine installed). It is also why a new sink is a function in one file +instead of another method on the executor. + **Tidy tables.** Parameters are `(dims…, value)`; a variable frame is `(dims…, var_label)`, one row per *existing* variable; a linear expression is `(frame dims…, var_label, coeff)` plus a constant part; constraint rows are @@ -274,7 +283,9 @@ than the plan level, which is | `errors.py` | the exception hierarchy; the one module the engine may import | | `relational/plan.py` | frozen logical-plan dataclasses | | `relational/arrow.py` | the Arrow boundary — caller tables in, via the PyCapsule protocol | -| `relational/executor.py` | duckdb execution + `lp_file` / `solver_direct` sinks | +| `relational/compiler.py` | plan → SQL text; pure, no connection | +| `relational/executor.py` | duckdb: bind sources, label, assemble the tables | +| `relational/sinks/` | how a built model leaves: `lp_file`, `solver_direct` (one module each, [README](linopy_yaml/relational/sinks/README.md)) | | `compat/__init__.py` | opt-in shim: `build` / `extend` on a `linopy.Model` | | `compat/loader.py` | data coercion to `xr.Dataset`, master coords | | `compat/builder.py` | eager backend: core AST → `linopy.Model` | diff --git a/linopy_yaml/relational/compiler.py b/linopy_yaml/relational/compiler.py new file mode 100644 index 00000000..67b07c5a --- /dev/null +++ b/linopy_yaml/relational/compiler.py @@ -0,0 +1,373 @@ +"""Logical plan → SQL. Pure: no connection, no execution, no I/O. + +`lowering.py` compiles the AST to a plan; this compiles the plan to SQL text. +Two stages, two names, and this is the second one. + +Nothing here runs anything. A :class:`SqlCompiler` needs three facts about the +model — the program, each dimension's cardinality, and which parameters are +boolean-valued — and from those it returns strings. That is what makes the +admissibility test in ARCHITECTURE.md ("read the verdict off the SQL") +something you can actually do: build a compiler, hand it a plan node, read the +SELECT it produces, and decide whether the operator is pointwise, bounded-halo +or global. No duckdb instance required, which is also why +``tests/test_compiler.py`` runs on a bare install. + +The unit of output is a :class:`TermFragment`: one additive piece of an affine +expression, carried as a full SELECT plus the dims it is indexed by. Compiling +an expression yields a term/const split, never a single query — because an LP +row *is* a sum of pieces, and keeping them separate is what lets every shape +operator rewrite one piece at a time. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from linopy_yaml.errors import LanguageError +from linopy_yaml.relational import plan + +if TYPE_CHECKING: + from collections.abc import Callable, Mapping + + +@dataclass(frozen=True) +class TermFragment: + """One additive piece of a compiled affine expression. + + ``sql`` is a full SELECT. Term fragments yield ``(dims…, var_label, + coeff)``; const fragments yield ``(dims…, cval)``. + """ + + dims: tuple[str, ...] + sql: str + is_term: bool + + +@dataclass(frozen=True) +class CompiledExpression: + """An affine expression as fragments: variable terms and a constant part.""" + + terms: tuple[TermFragment, ...] + consts: tuple[TermFragment, ...] + + +@dataclass(frozen=True) +class SqlCompiler: + """Turn plan nodes into SQL over the model's tidy tables. + + ``dimension_cardinality`` and ``boolean_parameters`` are read off the data + once it is loaded, which is the only reason this is not a free function: + ``sum`` over a dim the operand lacks scales by that dim's size, and + ``defined`` on a boolean parameter tests the value rather than its + finiteness. + """ + + program: plan.Program + dimension_cardinality: Mapping[str, int] + boolean_parameters: frozenset[str] + + # ------------------------------------------------------------------ + # frames — the masked coordinate product a declaration is instantiated over + # ------------------------------------------------------------------ + + def frame(self, dims: tuple[str, ...], where: plan.Predicate | None) -> tuple[str, str, str]: + """FROM/WHERE clauses of the (masked) coord product and its order key. + + Returns ``(from_clause, where_clause, order_key)``; the select list can + project ``t_.val AS ``. + """ + froms = [f'dim_{dims[0]} t_{dims[0]}'] + froms += [f'CROSS JOIN dim_{d} t_{d}' for d in dims[1:]] + + conds: list[str] = [] + if where is not None: + joins, cond = self.predicate(where, dims) + froms += joins + conds.append(cond) + return ' '.join(froms), (' AND '.join(conds) if conds else 'TRUE'), ', '.join(f't_{d}.ord' for d in dims) + + def parameter_join( + self, + param: str, + frame_dims: tuple[str, ...], + alias: str, + coordinate: str, + subject: str, + ) -> str: + """The ``LEFT JOIN`` clause binding *param* to the frame under *alias*. + + Both callers — where-masks and variable bounds — need the same join and + the same containment check: a parameter carrying a dim the frame does + not have would be reduced over that dim, silently widening a mask or + picking an arbitrary bound. They differ only in how the frame spells a + coordinate column (*coordinate*, a ``{dim}`` template) and in what the + message calls the offending parameter (*subject*) — naming the + declaration it came from is most of the value of raising here, so it is + the caller's word, not a role prefix pasted on the front. + """ + 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)}') + on = ' AND '.join(f'{alias}.{d} = {coordinate.format(dim=d)}' for d in declaration.dims) or 'TRUE' + return f'LEFT JOIN p_{param} {alias} ON {on}' + + # ------------------------------------------------------------------ + # predicates (where masks — row absence) + # ------------------------------------------------------------------ + + def predicate(self, pred: plan.Predicate, dims: tuple[str, ...]) -> tuple[list[str], str]: + """``(join clauses, boolean condition)`` for a where-mask over *dims*.""" + joins: dict[str, str] = {} + + def join_param(param: str) -> str: + alias = f'w_{param}' + joins[alias] = self.parameter_join(param, dims, alias, 't_{dim}.val', f"where-parameter '{param}'") + return alias + + def walk(p: plan.Predicate) -> str: + if isinstance(p, plan.ParameterComparison): + return _comparison_sql(f'{join_param(p.parameter)}.value', p.op, p.value) + if isinstance(p, plan.DimensionComparison): + if p.dimension not in dims: + raise LanguageError( + f"where-comparison on dimension '{p.dimension}' is outside the foreach dims " + f'{list(dims)} — reducing a mask over an unlisted dim is not supported' + ) + return _comparison_sql(f't_{p.dimension}.val', p.op, p.value) + if isinstance(p, plan.ParameterDefined): + alias = join_param(p.parameter) + if p.parameter in self.boolean_parameters: + return f'({alias}.value IS NOT NULL AND {alias}.value)' + return f'({alias}.value IS NOT NULL AND isfinite({alias}.value))' + if isinstance(p, plan.BooleanConstant): + return 'TRUE' if p.value else 'FALSE' + if isinstance(p, (plan.And, plan.Or)): + op = 'AND' if isinstance(p, plan.And) else 'OR' + return f'({walk(p.left)} {op} {walk(p.right)})' + if isinstance(p, plan.Not): + return f'(NOT COALESCE({walk(p.operand)}, FALSE))' + raise LanguageError(f'unsupported predicate node {type(p).__name__}') + + cond = walk(pred) + # NULL comparisons (missing parameter rows) must exclude the row, not + # yield NULL — wrap so the frame filter is strictly boolean. + return list(joins.values()), f'COALESCE({cond}, FALSE)' + + # ------------------------------------------------------------------ + # bounds + # ------------------------------------------------------------------ + + def bound(self, expr: plan.Expression, v: plan.VariableDeclaration) -> tuple[str, list[str]]: + """Compile a variable-free bound expression to a scalar SQL expression + over alias ``f`` (the variable frame), returning (sql, join clauses).""" + joins: dict[str, str] = {} + + def walk(e: plan.Expression) -> str: + if isinstance(e, plan.Constant): + return _literal(e.value) + if isinstance(e, plan.Parameter): + alias = f'b_{e.name}' + joins[alias] = self.parameter_join( + e.name, v.dims, alias, 'f.{dim}', f"bound parameter '{e.name}' of variable '{v.name}'" + ) + return f'{alias}.value' + if isinstance(e, plan.Negate): + return f'(-({walk(e.operand)}))' + if isinstance(e, plan.Add): + return f'({walk(e.left)} + {walk(e.right)})' + if isinstance(e, plan.Multiply): + return f'({walk(e.left)} * {walk(e.right)})' + raise LanguageError( + f"unsupported node {type(e).__name__} in bounds of variable '{v.name}' " + f'(bounds must be variable-free arithmetic over Constant/Parameter)' + ) + + return walk(expr), list(joins.values()) + + # ------------------------------------------------------------------ + # expressions → fragments + # ------------------------------------------------------------------ + + def expression(self, expr: plan.Expression, context: str) -> CompiledExpression: + """Compile an affine expression into term and const fragments.""" + + def ev(e: plan.Expression) -> CompiledExpression: + if isinstance(e, plan.Constant): + return CompiledExpression((), (TermFragment((), f'SELECT {_literal(e.value)} AS cval', False),)) + if isinstance(e, plan.Parameter): + d = self.program.parameter(e.name).dims + cols = ', '.join([*d, 'value AS cval']) if d else 'value AS cval' + return CompiledExpression((), (TermFragment(d, f'SELECT {cols} FROM p_{e.name}', False),)) + if isinstance(e, plan.Variable): + d = self.program.variable(e.name).dims + cols = ', '.join([*d, 'var_label', '1.0 AS coeff']) + return CompiledExpression((TermFragment(d, f'SELECT {cols} FROM var_{e.name}', True),), ()) + if isinstance(e, plan.Negate): + return _map_fragments(ev(e.operand), _negate) + if isinstance(e, plan.Add): + a, b = ev(e.left), ev(e.right) + return CompiledExpression(a.terms + b.terms, a.consts + b.consts) + if isinstance(e, plan.Multiply): + a, b = ev(e.left), ev(e.right) + if a.terms and b.terms: + raise LanguageError(f'nonlinear product in {context}: both factors contain variables') + if b.terms: # normalise: terms on the left + a, b = b, a + terms = tuple(_join_mul(t, c, is_term=True) for t in a.terms for c in b.consts) + consts = tuple(_join_mul(x, c, is_term=False) for x in a.consts for c in b.consts) + return CompiledExpression(terms, consts) + if isinstance(e, plan.Divide): + a, b = ev(e.numerator), ev(e.divisor) + if b.terms: + raise LanguageError(f'nonlinear quotient in {context}: the divisor contains variables') + if len(b.consts) != 1: + raise LanguageError( + f'in {context}: a divisor must be a single Constant/Parameter factor, ' + f'not a sum — rewrite as multiplication by a precomputed parameter' + ) + inv = b.consts[0] + terms = tuple(_join_mul(t, inv, is_term=True, op='/') for t in a.terms) + consts = tuple(_join_mul(x, inv, is_term=False, op='/') for x in a.consts) + return CompiledExpression(terms, consts) + if isinstance(e, plan.Sum): + return _map_fragments(ev(e.operand), lambda p: self._sum_fragment(p, e.over, context)) + if isinstance(e, plan.GroupSum): + return _map_fragments(ev(e.operand), lambda p: self._group_fragment(p, e, context)) + if isinstance(e, plan.Translate): + return _map_fragments(ev(e.operand), lambda p: self._translate_fragment(p, e, context)) + raise LanguageError(f'unsupported expression node {type(e).__name__} in {context}') + + return ev(expr) + + # ------------------------------------------------------------------ + # shape operators — one dim rewritten per fragment + # ------------------------------------------------------------------ + + def _sum_fragment(self, p: TermFragment, over: tuple[str, ...], context: str) -> TermFragment: + missing = [d for d in over if d not in p.dims] + if missing and not p.is_term: + raise LanguageError( + f'in {context}: Sum over {list(over)} of a constant part lacking dims ' + f'{missing} is ambiguous under masks — multiply explicitly instead' + ) + keep = tuple(d for d in p.dims if d not in over) + scale = math.prod(self.dimension_cardinality[d] for d in missing) + valcols = 'var_label, coeff' if p.is_term else 'cval' + if scale != 1: + valcols = f'var_label, coeff * {scale} AS coeff' if p.is_term else f'cval * {scale} AS cval' + cols = ', '.join([*keep, valcols]) if keep else valcols + return TermFragment(keep, f'SELECT {cols} FROM ({p.sql})', p.is_term) + + def _group_fragment(self, p: TermFragment, g: plan.GroupSum, context: str) -> TermFragment: + """Relabel dim ``over`` to ``into`` through a declared coordinate. + + No aggregate here: the fragment's dim tuple changes and duplicate + (row, col) pairs collapse in the terminal ``SUM(coeff)`` at assembly — + the same shape as :meth:`_sum_fragment` dropping a dim. The join is + against the dim table, whose coordinate column was checked for + containment at build time, so it cannot silently drop a term. + """ + if g.over not in p.dims: + raise LanguageError(f"in {context}: GroupSum over '{g.over}' but the expression has dims {list(p.dims)}") + keep = tuple(x for x in p.dims if x != g.over) + valcols = 't.var_label, t.coeff' if p.is_term else 't.cval' + keepcols = ', '.join([*(f't.{x}' for x in keep), f'g."{g.coordinate}" AS {g.into}', valcols]) + sql = f'SELECT {keepcols} FROM ({p.sql}) t JOIN dim_{g.over} g ON g.val = t.{g.over}' + return TermFragment((*keep, g.into), sql, p.is_term) + + def _translate_fragment(self, p: TermFragment, s: plan.Translate, context: str) -> TermFragment: + """Translation = a pointwise remap of the dim through its ord: + a row at ord *o* contributes to the output coord at ord ``(o + by) % + card``. No window function involved — bounded-halo locality.""" + if s.dimension not in p.dims: + raise LanguageError( + f"in {context}: translation along '{s.dimension}' but the expression has dims {list(p.dims)}" + ) + card = self.dimension_cardinality[s.dimension] + others = [d for d in p.dims if d != s.dimension] + valcols = 't.var_label, t.coeff' if p.is_term else 't.cval' + cols = ', '.join([*(f't.{d}' for d in others), f'd_out.val AS {s.dimension}', valcols]) + if s.wrap: + on = f'd_out.ord = ((d_in.ord + {s.by}) % {card} + {card}) % {card}' + else: + # acyclic: out-of-range rows simply don't join — zero contribution + on = f'd_out.ord = d_in.ord + {s.by}' + sql = ( + f'SELECT {cols} FROM ({p.sql}) t ' + f'JOIN dim_{s.dimension} d_in ON d_in.val = t.{s.dimension} ' + f'JOIN dim_{s.dimension} d_out ON {on}' + ) + return TermFragment(p.dims, sql, p.is_term) + + # ------------------------------------------------------------------ + # assembly helpers used by the executor + # ------------------------------------------------------------------ + + @staticmethod + def constant_scalar(p: TermFragment) -> str: + """Correlated scalar: the summed const fragment value for frame row ``f``.""" + cond = ' AND '.join(f'q.{d} = f.{d}' for d in p.dims) or 'TRUE' + return f'SELECT SUM(q.cval) FROM ({p.sql}) q WHERE {cond}' + + +def _literal(v: float) -> str: + if math.isinf(v): + return "('infinity'::DOUBLE)" if v > 0 else "('-infinity'::DOUBLE)" + # _literal(0) type-checks (int -> float) and would emit '0': INTEGER, not DOUBLE + # pyrefly: ignore[unnecessary-type-conversion] + return repr(float(v)) + + +def _map_fragments( + compiled: CompiledExpression, + rewrite: Callable[[TermFragment], TermFragment], +) -> CompiledExpression: + """Apply *rewrite* to every fragment, keeping the term/const split. + + Sum, GroupSum and Translate all rewrite each fragment on its own, which is + what pointwise and bounded-halo locality mean in code (ARCHITECTURE.md, + "Read the verdict off the SQL"). A node that needed the fragments + *together* would be a global operator, rejected at lowering instead. + """ + return CompiledExpression( + tuple(rewrite(p) for p in compiled.terms), + tuple(rewrite(p) for p in compiled.consts), + ) + + +def _comparison_sql(column: str, op: plan.ComparisonOperator, value: float | str) -> str: + """One where-comparison: ``( )``. + + The language's ``==`` is SQL's ``=``, and a string literal needs quoting — + stated once, since the parameter and dimension cases differ only in which + column they test. + """ + literal = f"'{value}'" if isinstance(value, str) else repr(value) + return f'({column} {"=" if op == "==" else op} {literal})' + + +def _negate(p: TermFragment) -> TermFragment: + cols = 'var_label, -coeff AS coeff' if p.is_term else '-cval AS cval' + sel = ', '.join([*p.dims, cols]) if p.dims else cols + return TermFragment(p.dims, f'SELECT {sel} FROM ({p.sql})', p.is_term) + + +def _join_mul(a: TermFragment, c: TermFragment, is_term: bool, op: str = '*') -> TermFragment: + """a op c where ``c`` is a const fragment; join on shared dims, broadcast the rest.""" + shared = [d for d in a.dims if d in c.dims] + on = ' AND '.join(f'a.{d} = c.{d}' for d in shared) or 'TRUE' + out_dims = a.dims + tuple(d for d in c.dims if d not in a.dims) + dimcols = [ + *(f'a.{d}' for d in a.dims), + *(f'c.{d}' for d in c.dims if d not in a.dims), + ] + val = f'a.var_label, a.coeff {op} c.cval AS coeff' if is_term else f'a.cval {op} c.cval AS cval' + sel = ', '.join([*dimcols, val]) + return TermFragment( + out_dims, + f'SELECT {sel} FROM ({a.sql}) a JOIN ({c.sql}) c ON {on}', + is_term, + ) diff --git a/linopy_yaml/relational/executor.py b/linopy_yaml/relational/executor.py index 4b29e55e..e9861d3d 100644 --- a/linopy_yaml/relational/executor.py +++ b/linopy_yaml/relational/executor.py @@ -1,21 +1,25 @@ -"""Duckdb executor for the logical plan. +"""Duckdb executor: fill the model tables, then hand them to a sink. The lane is described in ARCHITECTURE.md, "The relational lane". -Compiles a :class:`~linopy_yaml.relational.plan.Program` into tidy tables inside -a file-backed duckdb database under a hard ``memory_limit``, then streams the -model out through a sink: ``write_lp`` (portability / differential oracle) or -``solve`` (solver_direct — batched HiGHS ``addCols``/``addRows``; the full -model never exists in this process's memory). +This module owns the *connection* and the tables in it. It does not own the +SQL — :mod:`linopy_yaml.relational.compiler` turns plan nodes into strings and +never touches a connection — and it does not own the way a model leaves: +:mod:`linopy_yaml.relational.sinks` drains the tables into LP text or into +HiGHS, one module per sink. + +What is left here is what genuinely needs the database: binding sources, +building the dim tables, assigning labels, assembling ``cols``/``obj``/ +``rows``/``A``, and joining solution values back to coordinates. Hand-managed chunking exists in exactly two places, both forced by operators duckdb cannot spill: -1. Label assignment — a global ``ROW_NUMBER`` window materialises its whole - input, so labels are assigned per-chunk of the leading dim with a running - offset. This is one generic mechanism; every operator inherits it. -2. LP-text ``string_agg`` in ``write_lp`` — string aggregates don't spill, - and a fixed conservative chunk size costs nothing in the debugging sink. +1. Label assignment (here) — a global ``ROW_NUMBER`` window materialises its + whole input, so labels are assigned per-chunk of the leading dim with a + running offset. This is one generic mechanism; every operator inherits it. +2. LP-text ``string_agg`` (in the sink) — string aggregates don't spill, and a + fixed conservative chunk size costs nothing in the debugging sink. Everything else — joins, scaling, masks, and the numeric hash aggregates that assemble ``A`` — delegates to duckdb's own spilling under ``memory_limit``. @@ -23,12 +27,12 @@ (joins/masks/group_sum) and bounded-halo (roll: t±k, which still works under label chunking because terms join the *global* variable table) compose freely; genuinely global operators (running sums, normalisations) must be rejected at -lowering with a rewrite hint (e.g. running sum → state-variable recurrence). +lowering with a rewrite hint (e.g. running sum -> state-variable recurrence). -duckdb, pyarrow, numpy and highspy are imported lazily. Arrow is the only -in-memory table this module knows: sources arrive as ``pyarrow.Table`` (or a -parquet path) and results leave as ``pyarrow.Table``, so no dataframe library -is a dependency of the lane. ``lowering.tidy_sources`` is where a caller's +duckdb, pyarrow and numpy are imported lazily. Arrow is the only in-memory +table this module knows: sources arrive as ``pyarrow.Table`` (or a parquet +path) and results leave as ``pyarrow.Table``, so no dataframe library is a +dependency of the lane. ``sources.tidy_sources`` is where a caller's pandas/polars/xarray object is turned into one. """ @@ -45,18 +49,17 @@ from typing import TYPE_CHECKING, Any, Literal from linopy_yaml.errors import DataError, LanguageError, LinopyYamlError -from linopy_yaml.relational import plan +from linopy_yaml.relational import plan, sinks from linopy_yaml.relational.arrow import as_table +from linopy_yaml.relational.compiler import SqlCompiler if TYPE_CHECKING: - from collections.abc import Callable, Mapping + from collections.abc import Mapping import pandas as pd _IDENT = re.compile(r'^[A-Za-z_][A-Za-z0-9_]*$') -_COPY_OPTS = "(FORMAT csv, HEADER false, QUOTE '', ESCAPE '')" - #: Column an index table carries its row position in. The space makes it #: unrepresentable as a declared name (``_IDENT``), so it cannot collide with a #: dimension or coordinate the caller's table already has. @@ -71,25 +74,6 @@ RelationalBuildError = LinopyYamlError -@dataclass(frozen=True) -class _TermFragment: - """One additive piece of a compiled affine expression. - - ``sql`` is a full SELECT. Term pieces yield ``(dims…, var_label, coeff)``; - const pieces yield ``(dims…, cval)``. - """ - - dims: tuple[str, ...] - sql: str - is_term: bool - - -@dataclass(frozen=True) -class _CompiledExpression: - terms: tuple[_TermFragment, ...] - consts: tuple[_TermFragment, ...] - - @dataclass class Solution: """Solve result. ``primal(name)`` joins labels back to coords. @@ -176,7 +160,7 @@ def __exit__(self, *exc: object) -> Literal[False]: class DuckdbExecutor: - """Build and sink a :class:`Program` relationally under a memory budget.""" + """Build a :class:`Program` into tables under a memory budget, then sink it.""" def __init__( self, @@ -204,6 +188,7 @@ def __init__( self._finalizer = weakref.finalize(self, _release, self._con, self.workdir if self._own_workdir else None) self._program: plan.Program | None = None + self._compiler: SqlCompiler | None = None self._bool_params: set[str] = set() self._dim_card: dict[str, int] = {} self._n_cols = 0 @@ -224,6 +209,11 @@ def build(self, program: plan.Program, sources: Mapping[str, Any]) -> None: self._create_param_table(p, sources) self._create_dim_tables(program, sources) + # the compiler is built after the data, because two of its answers + # depend on it: sum over an absent dim scales by that dim's size, and + # `defined` on a boolean parameter tests the value, not its finiteness + self._compiler = SqlCompiler(program, dict(self._dim_card), frozenset(self._bool_params)) + self._con.execute('CREATE TABLE cols (col BIGINT, lb DOUBLE, ub DOUBLE, vtype VARCHAR)') self._con.execute('CREATE TABLE obj (col BIGINT, coeff DOUBLE)') self._con.execute('CREATE TABLE rows (row BIGINT, sense VARCHAR, rhs DOUBLE)') @@ -249,6 +239,11 @@ def _validate_names(self, program: plan.Program) -> None: if not _IDENT.match(n): raise LanguageError(f"name '{n}' is not a valid identifier ([A-Za-z_][A-Za-z0-9_]*)") + @property + def _sql(self) -> SqlCompiler: + assert self._compiler is not None, 'build() has not run' + return self._compiler + def _create_param_table(self, p: plan.ParameterDeclaration, sources: Mapping[str, Any]) -> None: if p.name not in sources: raise DataError(f"no source bound for parameter '{p.name}'") @@ -419,95 +414,6 @@ def _check_coordinate_containment(self, d: str, cname: str, target: str) -> None f'solves without them.' ) - # ------------------------------------------------------------------ - # frames (masked coord products with partition-wise labels) - # ------------------------------------------------------------------ - - def _frame_sql(self, dims: tuple[str, ...], where: plan.Predicate | None) -> tuple[str, str, str]: - """FROM/WHERE clauses of the (masked) coord product and its order key. - - Returns ``(from_clause, where_clause, order_key)``; the select list can - project ``t_.val AS ``. - """ - assert self._program is not None - froms = [f'dim_{dims[0]} t_{dims[0]}'] - froms += [f'CROSS JOIN dim_{d} t_{d}' for d in dims[1:]] - - conds: list[str] = [] - if where is not None: - joins, cond = self._pred_sql(where, dims) - froms += joins - conds.append(cond) - from_clause = ' '.join(froms) - where_clause = ' AND '.join(conds) if conds else 'TRUE' - order_key = ', '.join(f't_{d}.ord' for d in dims) - return from_clause, where_clause, order_key - - def _parameter_join( - self, - param: str, - frame_dims: tuple[str, ...], - alias: str, - coordinate: str, - subject: str, - ) -> str: - """The ``LEFT JOIN`` clause binding *param* to the frame under *alias*. - - Both callers — where-masks and variable bounds — need the same join and - the same containment check: a parameter carrying a dim the frame does - not have would be reduced over that dim, silently widening a mask or - picking an arbitrary bound. They differ only in how the frame spells a - coordinate column (*coordinate*, a ``{dim}`` template) and in what the - message calls the offending parameter (*subject*) — naming the - declaration it came from is most of the value of raising here, so it is - the caller's word, not a role prefix pasted on the front. - """ - assert self._program is not None - 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)}') - on = ' AND '.join(f'{alias}.{d} = {coordinate.format(dim=d)}' for d in declaration.dims) or 'TRUE' - return f'LEFT JOIN p_{param} {alias} ON {on}' - - def _pred_sql(self, pred: plan.Predicate, dims: tuple[str, ...]) -> tuple[list[str], str]: - assert self._program is not None - joins: dict[str, str] = {} - - def join_param(param: str) -> str: - alias = f'w_{param}' - joins[alias] = self._parameter_join(param, dims, alias, 't_{dim}.val', f"where-parameter '{param}'") - return alias - - def walk(p: plan.Predicate) -> str: - if isinstance(p, plan.ParameterComparison): - return _comparison_sql(f'{join_param(p.parameter)}.value', p.op, p.value) - if isinstance(p, plan.DimensionComparison): - if p.dimension not in dims: - raise LanguageError( - f"where-comparison on dimension '{p.dimension}' is outside the foreach dims " - f'{list(dims)} — reducing a mask over an unlisted dim is not supported' - ) - return _comparison_sql(f't_{p.dimension}.val', p.op, p.value) - if isinstance(p, plan.ParameterDefined): - alias = join_param(p.parameter) - if p.parameter in self._bool_params: - return f'({alias}.value IS NOT NULL AND {alias}.value)' - return f'({alias}.value IS NOT NULL AND isfinite({alias}.value))' - if isinstance(p, plan.BooleanConstant): - return 'TRUE' if p.value else 'FALSE' - if isinstance(p, (plan.And, plan.Or)): - op = 'AND' if isinstance(p, plan.And) else 'OR' - return f'({walk(p.left)} {op} {walk(p.right)})' - if isinstance(p, plan.Not): - return f'(NOT COALESCE({walk(p.operand)}, FALSE))' - raise LanguageError(f'unsupported predicate node {type(p).__name__}') - - cond = walk(pred) - # NULL comparisons (missing parameter rows) must exclude the row, not - # yield NULL — wrap so the frame filter is strictly boolean. - return list(joins.values()), f'COALESCE({cond}, FALSE)' - def _chunk_starts(self, lead_dim: str, other_card: float) -> list[tuple[int, int]]: card = self._dim_card[lead_dim] per_chunk = max(1, int(self.chunk_rows // max(1.0, other_card))) @@ -532,7 +438,7 @@ def _label_frame( disagree about which coordinate gets which solver index. """ collist = ', '.join(f't_{d}.val AS {d}' for d in dims) - from_clause, where_clause, order_key = self._frame_sql(dims, where) + from_clause, where_clause, order_key = self._sql.frame(dims, where) self._con.execute( f'CREATE TABLE {table} AS SELECT {collist}, 0::BIGINT AS {label} FROM {from_clause} WHERE FALSE' ) @@ -557,8 +463,8 @@ def _build_variable(self, v: plan.VariableDeclaration) -> None: raise LanguageError(f"variable '{v.name}' has no dims (scalars: use dims of size 1)") self._n_cols = self._label_frame(f'var_{v.name}', v.dims, v.where, 'var_label', self._n_cols) - lb_sql, lb_joins = self._bound_sql(v.lower, v) - ub_sql, ub_joins = self._bound_sql(v.upper, v) + lb_sql, lb_joins = self._sql.bound(v.lower, v) + ub_sql, ub_joins = self._sql.bound(v.upper, v) joins = ' '.join(dict.fromkeys(lb_joins + ub_joins)) self._con.execute( f"INSERT INTO cols SELECT f.var_label, {lb_sql}, {ub_sql}, '{v.variable_type}' FROM var_{v.name} f {joins}" @@ -570,155 +476,11 @@ def _build_variable(self, v: plan.VariableDeclaration) -> None: f'is missing values for some coordinates' ) - def _bound_sql(self, expr: plan.Expression, v: plan.VariableDeclaration) -> tuple[str, list[str]]: - """Compile a variable-free bound expression to a scalar SQL expression - over alias ``f`` (the variable frame), returning (sql, join clauses).""" - assert self._program is not None - joins: dict[str, str] = {} - - def walk(e: plan.Expression) -> str: - if isinstance(e, plan.Constant): - return _lit(e.value) - if isinstance(e, plan.Parameter): - alias = f'b_{e.name}' - joins[alias] = self._parameter_join( - e.name, v.dims, alias, 'f.{dim}', f"bound parameter '{e.name}' of variable '{v.name}'" - ) - return f'{alias}.value' - if isinstance(e, plan.Negate): - return f'(-({walk(e.operand)}))' - if isinstance(e, plan.Add): - return f'({walk(e.left)} + {walk(e.right)})' - if isinstance(e, plan.Multiply): - return f'({walk(e.left)} * {walk(e.right)})' - raise LanguageError( - f"unsupported node {type(e).__name__} in bounds of variable '{v.name}' " - f'(bounds must be variable-free arithmetic over Constant/Parameter)' - ) - - return walk(expr), list(joins.values()) - - # ------------------------------------------------------------------ - # expression compilation → pieces - # ------------------------------------------------------------------ - - def _compile(self, expr: plan.Expression, context: str) -> _CompiledExpression: - assert self._program is not None - prog = self._program - - def ev(e: plan.Expression) -> _CompiledExpression: - if isinstance(e, plan.Constant): - return _CompiledExpression((), (_TermFragment((), f'SELECT {_lit(e.value)} AS cval', False),)) - if isinstance(e, plan.Parameter): - d = prog.parameter(e.name).dims - cols = ', '.join([*d, 'value AS cval']) if d else 'value AS cval' - return _CompiledExpression((), (_TermFragment(d, f'SELECT {cols} FROM p_{e.name}', False),)) - if isinstance(e, plan.Variable): - d = prog.variable(e.name).dims - cols = ', '.join([*d, 'var_label', '1.0 AS coeff']) - return _CompiledExpression((_TermFragment(d, f'SELECT {cols} FROM var_{e.name}', True),), ()) - if isinstance(e, plan.Negate): - return _map_fragments(ev(e.operand), _negate) - if isinstance(e, plan.Add): - a, b = ev(e.left), ev(e.right) - return _CompiledExpression(a.terms + b.terms, a.consts + b.consts) - if isinstance(e, plan.Multiply): - a, b = ev(e.left), ev(e.right) - if a.terms and b.terms: - raise LanguageError(f'nonlinear product in {context}: both factors contain variables') - if b.terms: # normalise: terms on the left - a, b = b, a - terms = tuple(_join_mul(t, c, is_term=True) for t in a.terms for c in b.consts) - consts = tuple(_join_mul(x, c, is_term=False) for x in a.consts for c in b.consts) - return _CompiledExpression(terms, consts) - if isinstance(e, plan.Divide): - a, b = ev(e.numerator), ev(e.divisor) - if b.terms: - raise LanguageError(f'nonlinear quotient in {context}: the divisor contains variables') - if len(b.consts) != 1: - raise LanguageError( - f'in {context}: a divisor must be a single Constant/Parameter factor, ' - f'not a sum — rewrite as multiplication by a precomputed parameter' - ) - inv = b.consts[0] - terms = tuple(_join_mul(t, inv, is_term=True, op='/') for t in a.terms) - consts = tuple(_join_mul(x, inv, is_term=False, op='/') for x in a.consts) - return _CompiledExpression(terms, consts) - if isinstance(e, plan.Sum): - return _map_fragments(ev(e.operand), lambda p: self._sum_fragment(p, e.over, context)) - if isinstance(e, plan.GroupSum): - return _map_fragments(ev(e.operand), lambda p: self._group_fragment(p, e, context)) - if isinstance(e, plan.Translate): - return _map_fragments(ev(e.operand), lambda p: self._translate_fragment(p, e, context)) - raise LanguageError(f'unsupported expression node {type(e).__name__} in {context}') - - return ev(expr) - - def _sum_fragment(self, p: _TermFragment, over: tuple[str, ...], context: str) -> _TermFragment: - missing = [d for d in over if d not in p.dims] - if missing and not p.is_term: - raise LanguageError( - f'in {context}: Sum over {list(over)} of a constant part lacking dims ' - f'{missing} is ambiguous under masks — multiply explicitly instead' - ) - keep = tuple(d for d in p.dims if d not in over) - scale = math.prod(self._dim_card[d] for d in missing) - valcols = 'var_label, coeff' if p.is_term else 'cval' - if scale != 1: - valcols = f'var_label, coeff * {scale} AS coeff' if p.is_term else f'cval * {scale} AS cval' - cols = ', '.join([*keep, valcols]) if keep else valcols - return _TermFragment(keep, f'SELECT {cols} FROM ({p.sql})', p.is_term) - - def _group_fragment(self, p: _TermFragment, g: plan.GroupSum, context: str) -> _TermFragment: - """Relabel dim ``over`` to ``into`` through a declared coordinate. - - No aggregate here: the fragment's dim tuple changes and duplicate - (row, col) pairs collapse in the terminal ``SUM(coeff)`` at assembly — - the same shape as :meth:`_sum_fragment` dropping a dim. The join is - against the dim table, whose coordinate column was checked for - containment at build time, so it cannot silently drop a term. - """ - if g.over not in p.dims: - raise LanguageError(f"in {context}: GroupSum over '{g.over}' but the expression has dims {list(p.dims)}") - keep = tuple(x for x in p.dims if x != g.over) - valcols = 't.var_label, t.coeff' if p.is_term else 't.cval' - keepcols = ', '.join([*(f't.{x}' for x in keep), f'g."{g.coordinate}" AS {g.into}', valcols]) - sql = f'SELECT {keepcols} FROM ({p.sql}) t JOIN dim_{g.over} g ON g.val = t.{g.over}' - return _TermFragment((*keep, g.into), sql, p.is_term) - - def _translate_fragment(self, p: _TermFragment, s: plan.Translate, context: str) -> _TermFragment: - """Translation = a pointwise remap of the dim through its ord: - a row at ord *o* contributes to the output coord at ord ``(o + by) % - card``. No window function involved — bounded-halo locality.""" - if s.dimension not in p.dims: - raise LanguageError( - f"in {context}: translation along '{s.dimension}' but the expression has dims {list(p.dims)}" - ) - card = self._dim_card[s.dimension] - others = [d for d in p.dims if d != s.dimension] - valcols = 't.var_label, t.coeff' if p.is_term else 't.cval' - cols = ', '.join([*(f't.{d}' for d in others), f'd_out.val AS {s.dimension}', valcols]) - if s.wrap: - on = f'd_out.ord = ((d_in.ord + {s.by}) % {card} + {card}) % {card}' - else: - # acyclic: out-of-range rows simply don't join — zero contribution - on = f'd_out.ord = d_in.ord + {s.by}' - sql = ( - f'SELECT {cols} FROM ({p.sql}) t ' - f'JOIN dim_{s.dimension} d_in ON d_in.val = t.{s.dimension} ' - f'JOIN dim_{s.dimension} d_out ON {on}' - ) - return _TermFragment(p.dims, sql, p.is_term) - - # ------------------------------------------------------------------ - # constraints and objective - # ------------------------------------------------------------------ - def _build_constraint(self, c: plan.ConstraintDeclaration) -> None: if not c.dims: raise LanguageError(f"constraint '{c.name}' has no dims") - lhs = self._compile(c.lhs, f"constraint '{c.name}' lhs") - rhs = self._compile(c.rhs, f"constraint '{c.name}' rhs") + lhs = self._sql.expression(c.lhs, f"constraint '{c.name}' lhs") + rhs = self._sql.expression(c.rhs, f"constraint '{c.name}' rhs") # normalise: terms on the left (rhs terms negated), consts on the right terms = [(p, 1.0) for p in lhs.terms] + [(p, -1.0) for p in rhs.terms] consts = [(p, 1.0) for p in rhs.consts] + [(p, -1.0) for p in lhs.consts] @@ -732,7 +494,7 @@ def _build_constraint(self, c: plan.ConstraintDeclaration) -> None: self._n_rows = self._label_frame(f'con_{c.name}', c.dims, c.where, 'row', self._n_rows) - rhs_sql = ' + '.join(f'{sign} * COALESCE(({self._agg_const_join(p, c.dims)}), 0)' for p, sign in consts) or '0' + rhs_sql = ' + '.join(f'{sign} * COALESCE(({self._sql.constant_scalar(p)}), 0)' for p, sign in consts) or '0' term_selects = [] for p, sign in terms: @@ -752,13 +514,8 @@ def _build_constraint(self, c: plan.ConstraintDeclaration) -> None: ) self._con.execute(f"INSERT INTO rows WITH f AS ({frame}) SELECT f.row, '{c.sense}', {rhs_sql} FROM f") - def _agg_const_join(self, p: _TermFragment, frame_dims: tuple[str, ...]) -> str: - """Correlated scalar: the summed const piece value for frame row ``f``.""" - cond = ' AND '.join(f'q.{d} = f.{d}' for d in p.dims) or 'TRUE' - return f'SELECT SUM(q.cval) FROM ({p.sql}) q WHERE {cond}' - def _build_objective(self, o: plan.ObjectiveDeclaration) -> None: - comp = self._compile(o.expression, 'objective') + comp = self._sql.expression(o.expression, 'objective') for p in comp.consts: if p.dims: raise LanguageError( @@ -772,153 +529,27 @@ def _build_objective(self, o: plan.ObjectiveDeclaration) -> None: self._obj_sense = o.sense # ------------------------------------------------------------------ - # sink: LP file + # sinks — see relational/sinks/; the executor only supplies the tables # ------------------------------------------------------------------ - def write_lp(self, path: str | Path) -> None: - path = Path(path) - parts = self.workdir / 'lp_parts' - parts.mkdir(exist_ok=True) - - self._con.execute(f"COPY (SELECT printf('%+.17g x%d', coeff, col) FROM obj) TO '{parts / 'obj'}' {_COPY_OPTS}") - - nnz = self._scalar('SELECT count(*) FROM A') - avg = max(1, nnz // max(1, self._n_rows)) - per_chunk = max(1, self.chunk_rows // avg) - con_parts = [] - for i, lo in enumerate(range(0, max(self._n_rows, 1), per_chunk)): - hi = min(lo + per_chunk, self._n_rows) - if hi <= lo: - break - part = parts / f'cons.{i}' - con_parts.append(part) - self._con.execute( - f""" - COPY ( - SELECT printf('c%d:', r.row) || chr(10) - || COALESCE(string_agg(printf('%+.17g x%d', a.coeff, a.col), chr(10)), '+0 x0') - || chr(10) - || printf('%s %.17g', CASE r.sense WHEN '==' THEN '=' ELSE r.sense END, r.rhs) - FROM rows r LEFT JOIN A a USING (row) - WHERE r.row >= {lo} AND r.row < {hi} - GROUP BY r.row, r.sense, r.rhs - ) TO '{part}' {_COPY_OPTS} - """ - ) - - self._con.execute( - f""" - COPY ( - SELECT CASE WHEN lb = '-infinity'::DOUBLE THEN '-infinity' ELSE printf('%.17g', lb) END - || printf(' <= x%d <= ', col) - || CASE WHEN ub = 'infinity'::DOUBLE THEN '+infinity' ELSE printf('%.17g', ub) END - FROM cols - ) TO '{parts / 'bounds'}' {_COPY_OPTS} - """ + def _tables(self) -> sinks.ModelTables: + return sinks.ModelTables( + connection=self._con, + workdir=self.workdir, + chunk_rows=self.chunk_rows, + column_count=self._n_cols, + row_count=self._n_rows, + objective_sense=self._obj_sense, + objective_constant=self._obj_const, ) - integrality_sections = [] - for vtype, keyword in (('binary', 'binary'), ('integer', 'general')): - n = self._scalar(f"SELECT count(*) FROM cols WHERE vtype = '{vtype}'") - if n: - part = parts / keyword - integrality_sections.append((keyword, part)) - self._con.execute( - f"COPY (SELECT printf('x%d', col) FROM cols WHERE vtype = '{vtype}') TO '{part}' {_COPY_OPTS}" - ) - - sense = b'min' if self._obj_sense == 'min' else b'max' - with open(path, 'wb') as f: - f.write(sense + b'\n\nobj:\n') - if self._obj_const: - f.write(f'{self._obj_const:+.17g}\n'.encode()) - _cat(f, parts / 'obj') - f.write(b'\ns.t.\n\n') - for part in con_parts: - _cat(f, part) - f.write(b'\nbounds\n') - _cat(f, parts / 'bounds') - for keyword, part in integrality_sections: - f.write(f'\n{keyword}\n'.encode()) - _cat(f, part) - f.write(b'\nend\n') - shutil.rmtree(parts) - - # ------------------------------------------------------------------ - # sink: solver_direct (HiGHS) - # ------------------------------------------------------------------ + def write_lp(self, path: str | Path) -> None: + """Sink the built model to an LP file.""" + sinks.write_lp_file(self._tables(), path) def solve(self, batch_rows: int = 100_000) -> Solution: - import highspy - import numpy as np - - inf = highspy.kHighsInf - h = highspy.Highs() - h.setOptionValue('output_flag', False) - - empty_i = np.empty(0, dtype=np.int32) - empty_f = np.empty(0, dtype=np.float64) - reader = self._con.execute( - 'SELECT c.col, c.lb, c.ub, c.vtype, COALESCE(o.coeff, 0) AS cost ' - 'FROM cols c LEFT JOIN obj o USING (col) ORDER BY c.col' - ).to_arrow_reader(batch_rows) - for batch in reader: - d = batch.to_pydict() - lb = np.nan_to_num(np.asarray(d['lb'], dtype=np.float64), neginf=-inf, posinf=inf) - ub = np.nan_to_num(np.asarray(d['ub'], dtype=np.float64), neginf=-inf, posinf=inf) - cost = np.asarray(d['cost'], dtype=np.float64) - h.addCols(len(cost), cost, lb, ub, 0, empty_i, empty_i, empty_f) - vtype = np.asarray(d['vtype']) - noncont = np.flatnonzero(vtype != 'continuous') - if len(noncont): - cols_idx = np.asarray(d['col'], dtype=np.int32)[noncont] - integrality = np.full(len(noncont), int(highspy.HighsVarType.kInteger), dtype=np.uint8) - h.changeColsIntegrality(len(noncont), cols_idx, integrality) - - for lo in range(0, max(self._n_rows, 1), batch_rows): - hi = min(lo + batch_rows, self._n_rows) - if hi <= lo: - break - rows = self._con.execute( - f'SELECT row, sense, rhs FROM rows WHERE row >= {lo} AND row < {hi} ORDER BY row' - ).fetchnumpy() - a = self._con.execute( - f'SELECT row, col, coeff FROM A WHERE row >= {lo} AND row < {hi} ORDER BY row' - ).fetchnumpy() - rhs = np.asarray(rows['rhs'], dtype=np.float64) - sense = rows['sense'] - rlb = np.where(sense == '<=', -inf, rhs) - rub = np.where(sense == '>=', inf, rhs) - starts = np.searchsorted(np.asarray(a['row']), np.asarray(rows['row'])).astype(np.int32) - h.addRows( - len(rhs), - rlb, - rub, - len(a['col']), - starts, - np.asarray(a['col'], dtype=np.int32), - np.asarray(a['coeff'], dtype=np.float64), - ) - - if self._obj_sense == 'max': - h.changeObjectiveSense(highspy.ObjSense.kMaximize) - h.run() - - status = str(h.getModelStatus()).rsplit('.', 1)[-1].removeprefix('k') - objective = h.getInfo().objective_function_value + self._obj_const - - import pyarrow as pa - - sol = pa.table( - { - 'col': pa.array(np.arange(self._n_cols, dtype=np.int64)), - 'value': pa.array(np.asarray(h.getSolution().col_value, dtype=np.float64)), - } - ) - self._con.execute('DROP TABLE IF EXISTS sol') - self._con.register('sol_src', sol) - self._con.execute('CREATE TABLE sol AS SELECT * FROM sol_src') - self._con.unregister('sol_src') + """Sink the built model straight into HiGHS and solve it.""" + status, objective = sinks.solve_direct(self._tables(), batch_rows) return Solution(status=status, objective=objective, _executor=self) def _solution_sql(self, name: str) -> str: @@ -962,68 +593,3 @@ def _release(con: Any, workdir: Path | None) -> None: con.close() if workdir is not None: shutil.rmtree(workdir, ignore_errors=True) - - -def _lit(v: float) -> str: - if math.isinf(v): - return "('infinity'::DOUBLE)" if v > 0 else "('-infinity'::DOUBLE)" - # _lit(0) type-checks (int -> float) and would emit '0': INTEGER, not DOUBLE - # pyrefly: ignore[unnecessary-type-conversion] - return repr(float(v)) - - -def _map_fragments( - compiled: _CompiledExpression, - rewrite: Callable[[_TermFragment], _TermFragment], -) -> _CompiledExpression: - """Apply *rewrite* to every fragment, keeping the term/const split. - - Sum, GroupSum and Translate all rewrite each fragment on its own, which is - what pointwise and bounded-halo locality mean in code (ARCHITECTURE.md, - "Read the verdict off the SQL"). A node that needed the fragments - *together* would be a global operator, rejected at lowering instead. - """ - return _CompiledExpression( - tuple(rewrite(p) for p in compiled.terms), - tuple(rewrite(p) for p in compiled.consts), - ) - - -def _comparison_sql(column: str, op: plan.ComparisonOperator, value: float | str) -> str: - """One where-comparison: ``( )``. - - The language's ``==`` is SQL's ``=``, and a string literal needs quoting — - stated once, since the parameter and dimension cases differ only in which - column they test. - """ - literal = f"'{value}'" if isinstance(value, str) else repr(value) - return f'({column} {"=" if op == "==" else op} {literal})' - - -def _negate(p: _TermFragment) -> _TermFragment: - cols = 'var_label, -coeff AS coeff' if p.is_term else '-cval AS cval' - sel = ', '.join([*p.dims, cols]) if p.dims else cols - return _TermFragment(p.dims, f'SELECT {sel} FROM ({p.sql})', p.is_term) - - -def _join_mul(a: _TermFragment, c: _TermFragment, is_term: bool, op: str = '*') -> _TermFragment: - """a op c where ``c`` is a const piece; join on shared dims, broadcast the rest.""" - shared = [d for d in a.dims if d in c.dims] - on = ' AND '.join(f'a.{d} = c.{d}' for d in shared) or 'TRUE' - out_dims = a.dims + tuple(d for d in c.dims if d not in a.dims) - dimcols = [ - *(f'a.{d}' for d in a.dims), - *(f'c.{d}' for d in c.dims if d not in a.dims), - ] - val = f'a.var_label, a.coeff {op} c.cval AS coeff' if is_term else f'a.cval {op} c.cval AS cval' - sel = ', '.join([*dimcols, val]) - return _TermFragment( - out_dims, - f'SELECT {sel} FROM ({a.sql}) a JOIN ({c.sql}) c ON {on}', - is_term, - ) - - -def _cat(f: Any, part: Path) -> None: - with open(part, 'rb') as src: - shutil.copyfileobj(src, f) diff --git a/linopy_yaml/relational/sinks/README.md b/linopy_yaml/relational/sinks/README.md new file mode 100644 index 00000000..00891a08 --- /dev/null +++ b/linopy_yaml/relational/sinks/README.md @@ -0,0 +1,69 @@ +# Sinks + +How a built model leaves the engine. These are the two boxes downstream of the +executor in [ARCHITECTURE.md](../../../ARCHITECTURE.md)'s pipeline. + +| Module | Sink | Needs | +|---|---|---| +| `tables.py` | — | the contract every sink reads | +| `lp_file.py` | `lp_file` — LP text | nothing beyond duckdb | +| `highs.py` | `solver_direct` — COO batches → HiGHS | `highspy` | + +## The contract + +A sink takes a `ModelTables` and nothing else: a connection holding `cols` +(col, lb, ub, vtype), `obj` (col, coeff), `rows` (row, sense, rhs) and `A` +(row, col, coeff), plus the counts it chunks by and the objective's sense and +constant — those last two live outside the tables because a constant has no +column to attach to. + +A sink never learns how the tables were filled, and the executor never learns +how they are drained. That is the point: adding `mps` is a new module here, not +another method on `DuckdbExecutor`. + +## Adding a sink + +1. New module, named for the sink. Take `ModelTables`, return whatever the sink + naturally returns (`None` for a writer, `(status, objective)` for a solver). +2. Re-export it from `__init__.py`. +3. Thin delegation on `DuckdbExecutor` — three lines, no logic. +4. If it needs an optional dependency, import it **inside the function**. The + module boundary is the fence; the lazy import is what keeps importing this + package free for callers who will never use that sink. +5. Stream. Nothing here may materialise the model — hard rule 4. Aggregate + inside duckdb, or hand the solver batches. + +## Why one module per sink + +Because that is where the fences are. `highspy` is an optional dependency of +`solver_direct` alone, and a caller that only writes LP files should not import +it. Splitting by *kind* — all writers together, all solvers together — would +put two optional imports in one module and a function that branches on which +solver you meant. + +When `mps` lands it may well belong beside `lp_file` (both are chunked `COPY` +of `printf`'d rows into part files, concatenated bytewise; they would share +`_cat`, the chunking and the float formatting). That is a decision to take with +the code in hand, not now — a `text.py` holding one function today would be a +guess about a sink that does not exist. + +## When Track 4 lands + +[Track 4](../../../ROADMAP.md#track-4--sink-capabilities) gives each sink a +declared capability table so `check(model, sink=...)` can answer "will this +sink take it". Two notes for whoever writes it: + +- The capability table belongs next to the sink it describes, in that sink's + module. `__init__.py` collects them; it does not own them. +- Make the set of sinks **closed**, like `helpers.BUILTINS` — not an open + `register_sink()`. An installed plugin that can change the answer to + `ly.check(model, sink=...)` is hard rule 5's failure mode one level down. + +## Known issue + +Neither file sink emits in a stable order: several `COPY` statements have no +`ORDER BY` and `preserve_insertion_order=false` is set on the connection, so +two runs of the same model produce byte-different files with identical content +— [#109](https://github.com/FBumann/linopy-yaml/issues/109). `solver_direct` is +unaffected; every read it makes is explicitly ordered, because `searchsorted` +requires it. diff --git a/linopy_yaml/relational/sinks/__init__.py b/linopy_yaml/relational/sinks/__init__.py new file mode 100644 index 00000000..9eebec4d --- /dev/null +++ b/linopy_yaml/relational/sinks/__init__.py @@ -0,0 +1,22 @@ +"""Sinks: how a built model leaves the engine. + +ARCHITECTURE.md's pipeline draws two boxes downstream of the executor. This +package is those boxes — one module each, because that is where the fences +are: ``highspy`` is an optional dependency of ``solver_direct`` alone, and a +caller that only writes LP files should never import it. + +A sink reads :class:`ModelTables` and nothing else. Neither sink knows how the +tables were filled, and the executor does not know how they are drained, which +is what makes the planned ``mps`` sink a new module here rather than another +method on the executor. +""" + +from linopy_yaml.relational.sinks.highs import solve_direct +from linopy_yaml.relational.sinks.lp_file import write_lp_file +from linopy_yaml.relational.sinks.tables import ModelTables + +__all__ = [ + 'ModelTables', + 'solve_direct', + 'write_lp_file', +] diff --git a/linopy_yaml/relational/sinks/highs.py b/linopy_yaml/relational/sinks/highs.py new file mode 100644 index 00000000..cabdfc23 --- /dev/null +++ b/linopy_yaml/relational/sinks/highs.py @@ -0,0 +1,94 @@ +"""The ``solver_direct`` sink: COO batches straight into HiGHS. + +No float→text→parse round trip — that is the whole reason this exists beside +:mod:`~linopy_yaml.relational.sinks.lp_file`. Columns arrive as arrow batches, +rows as numpy slices of ``A``, and the full model never lands in one array. + +``highspy`` is imported inside the function rather than at module scope: it is +an optional dependency, and importing this module must stay free for callers +that only ever write LP files. The module boundary is the fence; the lazy +import is what keeps the fence cheap. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from linopy_yaml.relational.sinks.tables import ModelTables + + +def solve_direct(model: ModelTables, batch_rows: int = 100_000) -> tuple[str, float]: + """Stream the model into HiGHS and solve it. Returns ``(status, objective)``. + + Leaves the primal values in a ``sol`` table on the connection, so reading + results back stays a label join like every other read — the caller owns + the mapping from solver column index to coordinates. + """ + import highspy + import numpy as np + + con = model.connection + inf = highspy.kHighsInf + h = highspy.Highs() + h.setOptionValue('output_flag', False) + + empty_i = np.empty(0, dtype=np.int32) + empty_f = np.empty(0, dtype=np.float64) + reader = con.execute( + 'SELECT c.col, c.lb, c.ub, c.vtype, COALESCE(o.coeff, 0) AS cost ' + 'FROM cols c LEFT JOIN obj o USING (col) ORDER BY c.col' + ).to_arrow_reader(batch_rows) + for batch in reader: + d = batch.to_pydict() + lb = np.nan_to_num(np.asarray(d['lb'], dtype=np.float64), neginf=-inf, posinf=inf) + ub = np.nan_to_num(np.asarray(d['ub'], dtype=np.float64), neginf=-inf, posinf=inf) + cost = np.asarray(d['cost'], dtype=np.float64) + h.addCols(len(cost), cost, lb, ub, 0, empty_i, empty_i, empty_f) + variable_type = np.asarray(d['vtype']) + noncontinuous = np.flatnonzero(variable_type != 'continuous') + if len(noncontinuous): + cols_idx = np.asarray(d['col'], dtype=np.int32)[noncontinuous] + integrality = np.full(len(noncontinuous), int(highspy.HighsVarType.kInteger), dtype=np.uint8) + h.changeColsIntegrality(len(noncontinuous), cols_idx, integrality) + + for lo, hi in model.row_chunks(batch_rows): + rows = con.execute( + f'SELECT row, sense, rhs FROM rows WHERE row >= {lo} AND row < {hi} ORDER BY row' + ).fetchnumpy() + a = con.execute(f'SELECT row, col, coeff FROM A WHERE row >= {lo} AND row < {hi} ORDER BY row').fetchnumpy() + rhs = np.asarray(rows['rhs'], dtype=np.float64) + sense = rows['sense'] + rlb = np.where(sense == '<=', -inf, rhs) + rub = np.where(sense == '>=', inf, rhs) + starts = np.searchsorted(np.asarray(a['row']), np.asarray(rows['row'])).astype(np.int32) + h.addRows( + len(rhs), + rlb, + rub, + len(a['col']), + starts, + np.asarray(a['col'], dtype=np.int32), + np.asarray(a['coeff'], dtype=np.float64), + ) + + if model.objective_sense == 'max': + h.changeObjectiveSense(highspy.ObjSense.kMaximize) + h.run() + + status = str(h.getModelStatus()).rsplit('.', 1)[-1].removeprefix('k') + objective = h.getInfo().objective_function_value + model.objective_constant + + import pyarrow as pa + + primal = pa.table( + { + 'col': pa.array(np.arange(model.column_count, dtype=np.int64)), + 'value': pa.array(np.asarray(h.getSolution().col_value, dtype=np.float64)), + } + ) + con.execute('DROP TABLE IF EXISTS sol') + con.register('sol_src', primal) + con.execute('CREATE TABLE sol AS SELECT * FROM sol_src') + con.unregister('sol_src') + return status, objective diff --git a/linopy_yaml/relational/sinks/lp_file.py b/linopy_yaml/relational/sinks/lp_file.py new file mode 100644 index 00000000..cdbf73fa --- /dev/null +++ b/linopy_yaml/relational/sinks/lp_file.py @@ -0,0 +1,102 @@ +"""The ``lp_file`` sink: the model as LP text. + +Portability, debugging, and the differential oracle. Every section is produced +by a duckdb ``COPY`` into a part file and the parts are concatenated bytewise, +so the LP text never exists in this process's memory either — only the file +handle does. + +The one hand-managed chunk in the engine that is not label assignment lives +here: string aggregates do not spill, so the constraint text is emitted in +fixed row ranges. That costs nothing in a debugging sink. + +Block order is not stable run to run (#109); the content is. +""" + +from __future__ import annotations + +import shutil +from pathlib import Path +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from linopy_yaml.relational.sinks.tables import ModelTables + +#: Raw lines, no CSV quoting to undo. +_COPY_OPTS = "(FORMAT csv, HEADER false, QUOTE '', ESCAPE '')" + + +def write_lp_file(model: ModelTables, path: str | Path) -> None: + """Write the model as LP text. + + Every section is produced by a duckdb ``COPY`` into a part file, then the + parts are concatenated bytewise — so the LP text never exists in this + process's memory either, only the file handle does. + """ + path = Path(path) + parts = model.workdir / 'lp_parts' + parts.mkdir(exist_ok=True) + con = model.connection + + con.execute(f"COPY (SELECT printf('%+.17g x%d', coeff, col) FROM obj) TO '{parts / 'obj'}' {_COPY_OPTS}") + + nnz = model.scalar('SELECT count(*) FROM A') + avg = max(1, nnz // max(1, model.row_count)) + con_parts = [] + for i, (lo, hi) in enumerate(model.row_chunks(max(1, model.chunk_rows // avg))): + part = parts / f'cons.{i}' + con_parts.append(part) + con.execute( + f""" + COPY ( + SELECT printf('c%d:', r.row) || chr(10) + || COALESCE(string_agg(printf('%+.17g x%d', a.coeff, a.col), chr(10)), '+0 x0') + || chr(10) + || printf('%s %.17g', CASE r.sense WHEN '==' THEN '=' ELSE r.sense END, r.rhs) + FROM rows r LEFT JOIN A a USING (row) + WHERE r.row >= {lo} AND r.row < {hi} + GROUP BY r.row, r.sense, r.rhs + ) TO '{part}' {_COPY_OPTS} + """ + ) + + con.execute( + f""" + COPY ( + SELECT CASE WHEN lb = '-infinity'::DOUBLE THEN '-infinity' ELSE printf('%.17g', lb) END + || printf(' <= x%d <= ', col) + || CASE WHEN ub = 'infinity'::DOUBLE THEN '+infinity' ELSE printf('%.17g', ub) END + FROM cols + ) TO '{parts / 'bounds'}' {_COPY_OPTS} + """ + ) + + integrality_sections = [] + for variable_type, keyword in (('binary', 'binary'), ('integer', 'general')): + if model.scalar(f"SELECT count(*) FROM cols WHERE vtype = '{variable_type}'"): + part = parts / keyword + integrality_sections.append((keyword, part)) + con.execute( + f"COPY (SELECT printf('x%d', col) FROM cols WHERE vtype = '{variable_type}') TO '{part}' {_COPY_OPTS}" + ) + + sense = b'min' if model.objective_sense == 'min' else b'max' + with open(path, 'wb') as f: + f.write(sense + b'\n\nobj:\n') + if model.objective_constant: + f.write(f'{model.objective_constant:+.17g}\n'.encode()) + _cat(f, parts / 'obj') + f.write(b'\ns.t.\n\n') + for part in con_parts: + _cat(f, part) + f.write(b'\nbounds\n') + _cat(f, parts / 'bounds') + for keyword, part in integrality_sections: + f.write(f'\n{keyword}\n'.encode()) + _cat(f, part) + f.write(b'\nend\n') + shutil.rmtree(parts) + + +def _cat(f: Any, part: Path) -> None: + with open(part, 'rb') as src: + shutil.copyfileobj(src, f) diff --git a/linopy_yaml/relational/sinks/tables.py b/linopy_yaml/relational/sinks/tables.py new file mode 100644 index 00000000..5f07f7f8 --- /dev/null +++ b/linopy_yaml/relational/sinks/tables.py @@ -0,0 +1,48 @@ +"""What every sink reads, and nothing more. + +The contract between the executor and the sinks: four tables in a connection, +plus the handful of scalars a writer needs to size its own chunking. A sink +that needs a fifth thing states it here, where both sides can see it. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import Iterator + from pathlib import Path + + +@dataclass(frozen=True) +class ModelTables: + """The built model, as a sink sees it. + + ``connection`` holds four tables — ``cols`` (col, lb, ub, vtype), ``obj`` + (col, coeff), ``rows`` (row, sense, rhs) and ``A`` (row, col, coeff). The + scalars alongside are the ones a sink cannot cheaply recover: the counts + it chunks by, and the objective's sense and constant, which live outside + the tables because a constant has no column to attach to. + """ + + connection: Any + workdir: Path + chunk_rows: int + column_count: int + row_count: int + objective_sense: str + objective_constant: float + + def scalar(self, sql: str) -> Any: + row = self.connection.execute(sql).fetchone() + assert row is not None + return row[0] + + def row_chunks(self, per_chunk: int) -> Iterator[tuple[int, int]]: + """``(lo, hi)`` half-open row ranges covering the constraint matrix.""" + for lo in range(0, max(self.row_count, 1), per_chunk): + hi = min(lo + per_chunk, self.row_count) + if hi <= lo: + return + yield lo, hi diff --git a/tests/test_architecture.py b/tests/test_architecture.py index ad7b3453..710a723c 100644 --- a/tests/test_architecture.py +++ b/tests/test_architecture.py @@ -170,19 +170,20 @@ def test_expansion_has_no_mutable_module_state(): ) -def test_every_ir_expr_node_is_handled_by_the_executor(): - """Two-tier economy: a primitive is not done until the executor consumes - it. Grep-level drift alarm — the differential tests prove semantics.""" - import linopy_yaml.relational.plan as plan +def test_every_plan_node_is_handled_by_the_compiler(): + """Two-tier economy: a primitive is not done until the engine consumes it. - executor_src = (PKG / 'relational' / 'executor.py').read_text() - unhandled = [cls.__name__ for cls in plan.Expression.__subclasses__() if f'plan.{cls.__name__}' not in executor_src] - assert not unhandled, f'plan.Expression nodes unknown to the executor: {unhandled}' + The compiler is the consumer — it is the module that turns plan nodes into + SQL, so a node it does not mention has no relational meaning however much + the executor moves around it. Grep-level drift alarm; the differential + tests prove semantics. + """ + import linopy_yaml.relational.plan as plan - unhandled_pred = [ - cls.__name__ for cls in plan.Predicate.__subclasses__() if f'plan.{cls.__name__}' not in executor_src - ] - assert not unhandled_pred, f'plan.Predicate nodes unknown to the executor: {unhandled_pred}' + compiler_src = (PKG / 'relational' / 'compiler.py').read_text() + for base in (plan.Expression, plan.Predicate): + unhandled = [c.__name__ for c in base.__subclasses__() if f'plan.{c.__name__}' not in compiler_src] + assert not unhandled, f'plan.{base.__name__} nodes unknown to the compiler: {unhandled}' def test_both_lanes_implement_exactly_the_closed_helper_set(): @@ -213,9 +214,16 @@ def test_both_lanes_implement_exactly_the_closed_helper_set(): assert not missing, f'built-in helpers with no lowering case: {missing}' -def test_architecture_doc_mentions_every_module(): - """ARCHITECTURE.md's module map stays complete (its own first paragraph).""" - doc = (PKG.parent / 'ARCHITECTURE.md').read_text() +def test_every_module_is_documented_somewhere(): + """No module is undocumented — but the doc need not be ARCHITECTURE.md. + + A subpackage that grows a member per variant (one sink per module) would + push its whole membership list into the top-level map, which is the thing + that map exists *not* to be. A ``README.md`` beside the code counts + instead: it is what you read when you open the directory, and it stays + next to the thing it describes. + """ + architecture = (PKG.parent / 'ARCHITECTURE.md').read_text() missing = [] for path in _all_modules(): name = path.name @@ -223,9 +231,14 @@ def test_architecture_doc_mentions_every_module(): continue # private plumbing (_notes) needs no doc entry if name == '__init__.py': continue - if name not in doc: + local_readme = path.parent / 'README.md' + documented = name in architecture or (local_readme.exists() and name in local_readme.read_text()) + if not documented: missing.append(str(path.relative_to(PKG))) - assert not missing, f"modules absent from ARCHITECTURE.md's map: {missing}" + assert not missing, ( + f'undocumented modules: {missing} — add each to ARCHITECTURE.md, or to a ' + f'README.md in its own directory if it is one member of a family' + ) def test_every_schema_model_is_strict(): diff --git a/tests/test_compiler.py b/tests/test_compiler.py new file mode 100644 index 00000000..aa14db07 --- /dev/null +++ b/tests/test_compiler.py @@ -0,0 +1,185 @@ +"""The compiler is pure, and this file is the proof. + +No duckdb, no highspy, no data — a plan node goes in, SQL text comes out. +That is the seam the split bought: before it, checking what SQL an operator +emits meant building a model and solving it. + +These assertions are deliberately about *shape*, not exact text. They pin the +properties ARCHITECTURE.md's admissibility test reads off the SQL — which dim +columns survive, whether an aggregate or a window appears, whether a mask +becomes a join or a filter — and leave formatting free to change. +""" + +from __future__ import annotations + +import pytest + +from linopy_yaml.errors import LanguageError +from linopy_yaml.relational import plan +from linopy_yaml.relational.compiler import SqlCompiler + +PROGRAM = plan.Program( + parameters=( + plan.ParameterDeclaration('cost', ('generator',)), + plan.ParameterDeclaration('load', ('snapshot',)), + plan.ParameterDeclaration('available', ('generator',)), + ), + variables=(plan.VariableDeclaration('p', ('snapshot', 'generator')),), + constraints=(), + objective=plan.ObjectiveDeclaration('min', plan.Variable('p')), + dimensions=( + plan.DimensionDeclaration('snapshot'), + plan.DimensionDeclaration('generator', coordinates=(('bus', 'bus'),)), + plan.DimensionDeclaration('bus'), + ), +) + +CARDINALITY = {'snapshot': 24, 'generator': 3, 'bus': 2} + + +def compiler(boolean_parameters: frozenset[str] = frozenset()) -> SqlCompiler: + return SqlCompiler(PROGRAM, CARDINALITY, boolean_parameters) + + +# --------------------------------------------------------------------------- +# expressions +# --------------------------------------------------------------------------- + + +def test_a_variable_compiles_to_one_term_fragment_over_its_dims(): + compiled = compiler().expression(plan.Variable('p'), 'test') + assert len(compiled.terms) == 1 + assert not compiled.consts + fragment = compiled.terms[0] + assert fragment.dims == ('snapshot', 'generator') + assert fragment.is_term + assert 'FROM var_p' in fragment.sql + + +def test_a_parameter_is_a_constant_fragment_not_a_term(): + compiled = compiler().expression(plan.Parameter('cost'), 'test') + assert not compiled.terms + assert compiled.consts[0].dims == ('generator',) + assert not compiled.consts[0].is_term + + +def test_addition_concatenates_fragments_rather_than_joining(): + """An LP row is a sum of terms, so ``+`` needs no SQL at all.""" + compiled = compiler().expression(plan.Variable('p') + plan.Variable('p'), 'test') + assert len(compiled.terms) == 2 + + +def test_a_product_of_two_variable_carrying_factors_is_refused(): + with pytest.raises(LanguageError, match='nonlinear product'): + compiler().expression(plan.Multiply(plan.Variable('p'), plan.Variable('p')), 'test') + + +def test_a_divisor_carrying_variables_is_refused(): + with pytest.raises(LanguageError, match='nonlinear quotient'): + compiler().expression(plan.Divide(plan.Variable('p'), plan.Variable('p')), 'test') + + +# --------------------------------------------------------------------------- +# shape operators — each rewrites exactly one dim column +# --------------------------------------------------------------------------- + + +def test_sum_drops_the_dim_it_sums_over(): + compiled = compiler().expression(plan.Sum(plan.Variable('p'), ('generator',)), 'test') + assert compiled.terms[0].dims == ('snapshot',) + + +def test_sum_over_an_absent_dim_scales_by_that_dims_cardinality(): + """Eager parity: summing a snapshot-only term over `generator` repeats it.""" + inner = plan.Sum(plan.Variable('p'), ('generator',)) + compiled = compiler().expression(plan.Sum(inner, ('generator',)), 'test') + assert 'coeff * 3' in compiled.terms[0].sql + + +def test_group_sum_swaps_the_source_dim_for_the_target_and_emits_no_aggregate(): + """The GROUP BY lives in the terminal assembly, not in the fragment — + which is what keeps the operator pointwise.""" + node = plan.GroupSum(plan.Variable('p'), over='generator', coordinate='bus', into='bus') + fragment = compiler().expression(node, 'test').terms[0] + assert fragment.dims == ('snapshot', 'bus') + assert 'GROUP BY' not in fragment.sql + assert 'JOIN dim_generator' in fragment.sql + + +def test_translate_keeps_its_dims_and_joins_the_dim_table_twice(): + """Bounded halo: a row at ord *o* lands at ord *o + by*, no window.""" + fragment = compiler().expression(plan.Translate(plan.Variable('p'), 'snapshot', by=1), 'test').terms[0] + assert fragment.dims == ('snapshot', 'generator') + assert fragment.sql.count('JOIN dim_snapshot') == 2 + assert 'OVER (' not in fragment.sql + + +def test_wrapping_is_modulo_and_acyclic_is_not(): + cyclic = compiler().expression(plan.Translate(plan.Variable('p'), 'snapshot', by=1, wrap=True), 't').terms[0] + acyclic = compiler().expression(plan.Translate(plan.Variable('p'), 'snapshot', by=1, wrap=False), 't').terms[0] + assert '% 24' in cyclic.sql + assert '%' not in acyclic.sql + + +def test_a_shape_operator_along_a_dim_the_expression_lacks_is_refused(): + with pytest.raises(LanguageError, match='translation'): + compiler().expression(plan.Translate(plan.Parameter('cost'), 'snapshot', by=1), 'test') + + +# --------------------------------------------------------------------------- +# predicates +# --------------------------------------------------------------------------- + + +def test_a_dimension_comparison_filters_a_column_already_in_the_frame(): + joins, condition = compiler().predicate(plan.DimensionComparison('snapshot', '>', 0), ('snapshot',)) + assert joins == [] + assert 't_snapshot.val > 0' in condition + + +def test_a_parameter_predicate_needs_a_left_join(): + joins, condition = compiler().predicate(plan.ParameterDefined('available'), ('generator',)) + assert len(joins) == 1 + assert 'LEFT JOIN p_available' in joins[0] + assert 'isfinite' in condition + + +def test_defined_on_a_boolean_parameter_tests_the_value_not_its_finiteness(): + _, condition = compiler(frozenset({'available'})).predicate(plan.ParameterDefined('available'), ('generator',)) + assert 'isfinite' not in condition + + +def test_a_mask_is_wrapped_so_a_null_excludes_the_row(): + _, condition = compiler().predicate(plan.ParameterComparison('available', '>', 0), ('generator',)) + assert condition.startswith('COALESCE(') + assert condition.endswith(', FALSE)') + + +def test_a_where_parameter_outside_the_frame_dims_is_refused(): + """Otherwise the mask would be reduced over a dim the declaration never named.""" + with pytest.raises(LanguageError, match='outside the foreach dims'): + compiler().predicate(plan.ParameterDefined('load'), ('generator',)) + + +# --------------------------------------------------------------------------- +# frames and bounds +# --------------------------------------------------------------------------- + + +def test_a_frame_cross_joins_its_dim_tables_and_orders_by_ordinal(): + from_clause, where_clause, order_key = compiler().frame(('snapshot', 'generator'), None) + assert from_clause == 'dim_snapshot t_snapshot CROSS JOIN dim_generator t_generator' + assert where_clause == 'TRUE' + assert order_key == 't_snapshot.ord, t_generator.ord' + + +def test_a_parameter_bound_joins_on_the_variable_frame(): + variable = PROGRAM.variables[0] + sql, joins = compiler().bound(plan.Parameter('cost'), variable) + assert sql == 'b_cost.value' + assert 'LEFT JOIN p_cost b_cost ON b_cost.generator = f.generator' in joins[0] + + +def test_a_bound_carrying_a_variable_is_refused(): + with pytest.raises(LanguageError, match='bounds must be variable-free'): + compiler().bound(plan.Variable('p'), PROGRAM.variables[0])