feat(codegen): Substrait front-end for the codegen serving engine

Add CodegenSubstraitFn, a second codegen backend that consumes a Substrait
plan instead of SQL. The plan is translated into the existing _codegen plan IR
and handed to the same backend-agnostic pipeline (optimize/validate/infer/
compile) and rt.* runtime — only the front-end differs.

- _codegen_substrait/: Context (row schemas + static tables), consumer
  (Substrait -> _codegen IR: Read/Filter/Project, scalar exprs, positional
  field refs, function anchors, literals, casts), and CodegenSubstraitFn.
- Reuse: extract CodegenFn._finalize so both the SQL and Substrait front-ends
  join at the plan IR; nothing downstream changes.
- Input is a Substrait plan (bytes); from_sql() produces one via DataFusion for
  callers/tests. Table names are resolved case-insensitively (DataFusion folds
  them). Surfaces DataFusion can't serialize (e.g. UNNEST) defer cleanly.
- Wire a third "substrait" backend into the differential harness; it passes the
  supported scalar/scan/filter/project surface and skips the rest.

Scope: the basic relational/scalar surface. Joins, string builtins, CASE and
transformers are deferred (skip, not fail).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: ahrzb

pyproject.toml

@@ -9,6 +9,7 @@ dependencies = [
     "pyarrow>=19.0",
     "pydantic>=2.0,<3.0",
     "sqlglot>=30.12.0",
+    "substrait>=0.29.0",
 ]
 
 [dependency-groups]

sql_transform/_codegen/engine.py

@@ -227,6 +227,19 @@ def __init__(
         output_model: type[BaseModel] | None = None,
     ) -> None:
         plan = cp.build_plan(sql)
+        self._finalize(plan, row_tables, static_tables, output_model)
+
+    def _finalize(
+        self,
+        plan: cp.Plan,
+        row_tables: dict,
+        static_tables: dict,
+        output_model: type[BaseModel] | None,
+    ) -> None:
+        """Shared pipeline from the plan IR onward — optimize, validate, infer
+        types, build the output model + static lookups, compile. Both the SQL
+        front-end (CodegenFn) and the Substrait front-end (CodegenSubstraitFn) join
+        here; everything downstream of the plan IR is backend-agnostic."""
         plan, specs = cp.optimize(plan, set(static_tables))
 
         row_schemas = {n: cp.schema_from_pydantic(m) for n, m in row_tables.items()}

sql_transform/_codegen_substrait/__init__.py

@@ -0,0 +1,20 @@
+"""Substrait front-end for the codegen serving engine.
+
+Same idea as `sql_transform._codegen`, but the plan comes from a **Substrait
+plan** instead of SQL parsed by sqlglot. The Substrait plan is translated into
+the existing `_codegen` plan IR and handed to the same backend-agnostic pipeline
+(`CodegenFn._finalize`): optimize, validate, type-infer, compile, run. The
+`rt.*` runtime and the emitter are reused untouched.
+
+References the Substrait plan against a `Context` (registered row-table schemas
+and static tables) — the plan carries only names/anchors/indices; the context
+resolves them.
+"""
+
+from __future__ import annotations
+
+from sql_transform._codegen.plan import UnsupportedInCodegen
+from sql_transform._codegen_substrait.context import Context
+from sql_transform._codegen_substrait.engine import CodegenSubstraitFn
+
+__all__ = ["CodegenSubstraitFn", "Context", "UnsupportedInCodegen"]

sql_transform/_codegen_substrait/consumer.py

@@ -0,0 +1,187 @@
+"""Translate a Substrait plan into the existing `_codegen` plan IR.
+
+Scope (the "basic part"): Read, Filter and the top-level Project, over scalar
+expressions — field references, literals, arithmetic/comparison/boolean scalar
+functions, a few string builtins, and casts. Anything else raises
+`UnsupportedInCodegen`, so the differential harness reports it as a deferred
+surface rather than a wrong answer.
+
+Field references in Substrait are POSITIONAL (an index into the emitting rel's
+output). We thread an ordered list of `(table, column_name)` down the tree so
+each index resolves back to the `Column(table, name)` the IR/emitter expect.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from substrait.proto import Plan
+
+from sql_transform._codegen import plan as cp
+
+# Substrait standard-function name -> IR binary-op name (matches plan._BINOPS).
+_BINOP = {
+    "add": "add",
+    "subtract": "sub",
+    "multiply": "mul",
+    "divide": "div",
+    "modulus": "mod",
+    "equal": "eq",
+    "not_equal": "neq",
+    "lt": "lt",
+    "gt": "gt",
+    "lte": "lte",
+    "gte": "gte",
+    "and": "and",
+    "or": "or",
+}
+# Substrait scalar function -> IR builtin name (matches engine._BUILTINS).
+_FUNC = {
+    "upper": "upper",
+    "lower": "lower",
+    "abs": "abs",
+    "coalesce": "coalesce",
+}
+# Substrait cast target type field -> IR base type.
+_CAST = {
+    "string": cp.STR,
+    "i64": cp.INT,
+    "i32": cp.INT,
+    "fp64": cp.FLOAT,
+    "fp32": cp.FLOAT,
+    "bool": cp.BOOL,
+}
+
+
+def consume(plan_bytes: bytes, table_names: Any = ()) -> cp.Plan:
+    """Substrait plan bytes -> `cp.Plan`. Column names come from the plan's own
+    Read schemas. `table_names` are the registered names (from the Context): a
+    Substrait Read folds its table name to lowercase (DataFusion catalog
+    behaviour), so we resolve it back case-insensitively."""
+    sp = Plan()
+    sp.ParseFromString(plan_bytes)
+
+    funcs: dict[int, str] = {}
+    for ext in sp.extensions:
+        if ext.HasField("extension_function"):
+            funcs[ext.extension_function.function_anchor] = ext.extension_function.name
+
+    if not sp.relations:
+        raise cp.UnsupportedInCodegen("Substrait plan has no relations")
+    root = sp.relations[0].root
+    aliases = list(root.names)
+    tmap = {n.lower(): n for n in table_names}
+    return _project(root.input, aliases, funcs, tmap)
+
+
+def _project(
+    rel: Any, aliases: list[str], funcs: dict[int, str], tmap: dict[str, str]
+) -> cp.Plan:
+    """The top rel must be a Project; peel it into `Plan(projection, input)`."""
+    if not rel.HasField("project"):
+        raise cp.UnsupportedInCodegen(
+            f"expected a top-level Project, got {rel.WhichOneof('rel_type')}"
+        )
+    proj = rel.project
+    input_rel, cols = _rel(proj.input, funcs, tmap)
+
+    # Substrait Project APPENDS its computed expressions after the input columns;
+    # `emit.output_mapping` (when present) selects which of that pool is output.
+    pool: list[Any] = [cp.Column(t, n) for t, n in cols]
+    pool += [_expr(e, cols, funcs) for e in proj.expressions]
+
+    mapping = (
+        list(proj.common.emit.output_mapping)
+        if proj.common.HasField("emit")
+        else list(range(len(pool)))
+    )
+    if aliases and len(aliases) != len(mapping):
+        raise cp.UnsupportedInCodegen("Substrait output name/column count mismatch")
+
+    projection = []
+    for i, idx in enumerate(mapping):
+        alias = aliases[i] if aliases else _col_name(pool[idx], i)
+        projection.append((alias, pool[idx]))
+    return cp.Plan(projection, input_rel)
+
+
+def _col_name(expr: Any, i: int) -> str:
+    return expr.name if isinstance(expr, cp.Column) else f"col{i}"
+
+
+def _rel(
+    rel: Any, funcs: dict[int, str], tmap: dict[str, str]
+) -> tuple[Any, list[tuple[str, str]]]:
+    """Translate a relational node, returning (IR rel, ordered [(table, col)])."""
+    kind = rel.WhichOneof("rel_type")
+    if kind == "read":
+        return _read(rel.read, tmap)
+    if kind == "filter":
+        inner, cols = _rel(rel.filter.input, funcs, tmap)
+        pred = _expr(rel.filter.condition, cols, funcs)
+        return cp.Filter(inner, pred), cols
+    raise cp.UnsupportedInCodegen(f"Substrait relation '{kind}' is not supported yet")
+
+
+def _read(read: Any, tmap: dict[str, str]) -> tuple[Any, list[tuple[str, str]]]:
+    if not read.HasField("named_table"):
+        raise cp.UnsupportedInCodegen("only named-table reads are supported")
+    folded = read.named_table.names[-1]
+    table = tmap.get(folded.lower(), folded)
+    names = list(read.base_schema.names)
+    # A projection selects a column subset (by index into base_schema), in order.
+    if read.HasField("projection"):
+        picked = [it.field for it in read.projection.select.struct_items]
+        names = [names[i] for i in picked]
+    return cp.TableScan(table), [(table, n) for n in names]
+
+
+def _expr(e: Any, cols: list[tuple[str, str]], funcs: dict[int, str]) -> Any:
+    kind = e.WhichOneof("rex_type")
+    if kind == "selection":
+        idx = e.selection.direct_reference.struct_field.field
+        table, name = cols[idx]
+        return cp.Column(table, name)
+    if kind == "literal":
+        return cp.Literal(_literal(e.literal))
+    if kind == "cast":
+        target = e.cast.type.WhichOneof("kind")
+        if target not in _CAST:
+            raise cp.UnsupportedInCodegen(f"cast to Substrait type '{target}'")
+        return cp.Cast(_expr(e.cast.input, cols, funcs), _CAST[target])
+    if kind == "scalar_function":
+        return _scalar(e.scalar_function, cols, funcs)
+    raise cp.UnsupportedInCodegen(f"Substrait expression '{kind}' is not supported yet")
+
+
+def _scalar(sf: Any, cols: list[tuple[str, str]], funcs: dict[int, str]) -> Any:
+    name = funcs.get(sf.function_reference)
+    if name is None:
+        raise cp.UnsupportedInCodegen(
+            f"unknown Substrait function anchor {sf.function_reference}"
+        )
+    args = [_expr(a.value, cols, funcs) for a in sf.arguments if a.HasField("value")]
+    if name == "not":
+        return cp.Not(args[0])
+    if name in _BINOP:
+        if len(args) != 2:
+            raise cp.UnsupportedInCodegen(f"{name} with {len(args)} args")
+        return cp.BinaryOp(_BINOP[name], args[0], args[1])
+    if name in _FUNC:
+        return cp.Func(_FUNC[name], args)
+    raise cp.UnsupportedInCodegen(f"Substrait function '{name}' is not supported yet")
+
+
+def _literal(lit: Any) -> Any:
+    kind = lit.WhichOneof("literal_type")
+    if kind == "null":
+        return None
+    if kind in ("fp64", "fp32"):
+        return float(getattr(lit, kind))
+    if kind in ("i64", "i32", "i16", "i8"):
+        return int(getattr(lit, kind))
+    if kind == "boolean":
+        return bool(lit.boolean)
+    if kind == "string":
+        return lit.string
+    raise cp.UnsupportedInCodegen(f"Substrait literal '{kind}' is not supported yet")

sql_transform/_codegen_substrait/context.py

@@ -0,0 +1,24 @@
+"""The resolution environment for a Substrait plan.
+
+A Substrait plan is portable: it names tables, references functions by anchor,
+and columns by position. It carries no schemas or table data. The `Context`
+supplies what the plan omits — the row-table schemas (as Pydantic models, the
+same shape `CodegenFn`/`InferFn` take) and the static tables (Arrow tables the
+lookup joins index). Registered in the DataFusion session that *produces* the
+plan and consulted again when we *consume* it.
+
+Transformers are intentionally out of scope for now (see package docstring).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass(frozen=True)
+class Context:
+    row_tables: dict[str, Any] = field(default_factory=dict)
+    """name -> Pydantic model class (the row schema)."""
+    static_tables: dict[str, Any] = field(default_factory=dict)
+    """name -> pyarrow.Table (frozen fit-state / lookup tables)."""

sql_transform/_codegen_substrait/engine.py

@@ -0,0 +1,85 @@
+"""`CodegenSubstraitFn` — a codegen serving engine driven by a Substrait plan.
+
+Takes a `Context` (registered row-table schemas + static tables) and a Substrait
+plan; translates the plan into the `_codegen` IR (via `consumer.consume`) and
+reuses `CodegenFn`'s backend-agnostic pipeline and `infer()` unchanged.
+
+`from_sql` is the conditional the caller uses to pass SQL instead of a plan: it
+produces the Substrait plan through DataFusion (registering the context's tables)
+and feeds it to the same consumer — so SQL and a hand-supplied plan are one path.
+"""
+
+from __future__ import annotations
+
+import datafusion
+import pyarrow as pa
+from datafusion import substrait as dss
+from pydantic import BaseModel
+
+from sql_transform._codegen import plan as cp
+from sql_transform._codegen.engine import CodegenFn
+from sql_transform._codegen_substrait.consumer import consume
+from sql_transform._codegen_substrait.context import Context
+
+_ARROW = {
+    cp.INT: pa.int64(),
+    cp.FLOAT: pa.float64(),
+    cp.STR: pa.string(),
+    cp.BOOL: pa.bool_(),
+}
+
+
+class CodegenSubstraitFn(CodegenFn):
+    """Codegen engine whose plan comes from Substrait, not SQL. Same `infer()`."""
+
+    def __init__(
+        self,
+        context: Context,
+        substrait_plan: bytes,
+        output_model: type[BaseModel] | None = None,
+    ) -> None:
+        table_names = set(context.row_tables) | set(context.static_tables)
+        plan = consume(substrait_plan, table_names)
+        self._finalize(plan, context.row_tables, context.static_tables, output_model)
+
+    @classmethod
+    def from_sql(
+        cls,
+        sql: str,
+        context: Context,
+        output_model: type[BaseModel] | None = None,
+    ) -> CodegenSubstraitFn:
+        """Produce the Substrait plan from SQL via DataFusion, then consume it."""
+        return cls(context, sql_to_substrait(sql, context), output_model)
+
+
+def sql_to_substrait(sql: str, context: Context) -> bytes:
+    """Serialize `sql` to a Substrait plan, registering the context's tables in a
+    fresh DataFusion session so column names/types and table names resolve."""
+    ctx = datafusion.SessionContext()
+    for name, model in context.row_tables.items():
+        empty = pa.RecordBatch.from_pylist([], schema=_arrow_schema(model))
+        ctx.register_record_batches(name, [[empty]])
+    for name, table in context.static_tables.items():
+        ctx.register_record_batches(name, [table.to_batches()])
+    try:
+        return dss.Serde.serialize_bytes(sql, ctx)
+    except Exception as e:  # noqa: BLE001 -- DataFusion raises a bare Exception
+        # DataFusion's own Substrait producer can't express this query (e.g.
+        # UNNEST). That surface is unreachable through Substrait — defer it so the
+        # differential harness skips rather than fails, matching codegen's pattern.
+        raise cp.UnsupportedInCodegen(
+            f"DataFusion cannot serialize this query to Substrait: {str(e)[:120]}"
+        ) from e
+
+
+def _arrow_schema(model: type[BaseModel]) -> pa.Schema:
+    fields = []
+    for name, ft in cp.schema_from_pydantic(model).items():
+        arrow_t = _ARROW.get(ft.base)
+        if arrow_t is None:
+            raise cp.UnsupportedInCodegen(
+                f"cannot register column '{name}': non-scalar type not supported yet"
+            )
+        fields.append(pa.field(name, arrow_t, nullable=ft.nullable))
+    return pa.schema(fields)