Refactor dataflow operators around runtime registry

src/govoplan_dataflow/backend/executor.py

@@ -4,9 +4,9 @@ import hashlib
 import json
 import time
 from collections import defaultdict
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from decimal import Decimal
-from typing import Any, Callable
+from typing import Any, Callable, NoReturn
 
 from govoplan_dataflow.backend.expressions import (
     convert_value,
@@ -84,6 +84,26 @@ class ResolvedSource:
 SourceResolver = Callable[[GraphNode, int], ResolvedSource]
 
 
+@dataclass(frozen=True, slots=True)
+class _PreviewPlan:
+    graph: PipelineGraph
+    node_by_id: dict[str, GraphNode]
+    inputs_by_port: dict[str, dict[str, list[str]]]
+    ordered_node_ids: list[str]
+    validation: list[DataflowDiagnostic]
+    preview_node_id: str | None
+    row_limit: int
+
+
+@dataclass(slots=True)
+class _PreviewState:
+    started: float = field(default_factory=time.monotonic)
+    outputs: dict[str, list[dict[str, Any]]] = field(default_factory=dict)
+    node_diagnostics: list[NodePreviewDiagnostic] = field(default_factory=list)
+    source_fingerprints: list[dict[str, Any]] = field(default_factory=list)
+    input_row_count: int = 0
+
+
 def execute_preview(
     graph: PipelineGraph,
     *,
@@ -92,151 +112,315 @@ def execute_preview(
     preview_node_id: str | None = None,
     _execution_depth: int = 0,
 ) -> PipelineExecutionResult:
-    if _execution_depth > 5:
+    plan = _prepare_preview(
+        graph,
+        row_limit=row_limit,
+        preview_node_id=preview_node_id,
+        execution_depth=_execution_depth,
+    )
+    state = _PreviewState()
+    for node_id in plan.ordered_node_ids:
+        _execute_preview_node(
+            plan,
+            state,
+            plan.node_by_id[node_id],
+            source_resolver=source_resolver,
+            execution_depth=_execution_depth,
+        )
+    return _preview_result(plan, state)
+
+
+def _prepare_preview(
+    graph: PipelineGraph,
+    *,
+    row_limit: int,
+    preview_node_id: str | None,
+    execution_depth: int,
+) -> _PreviewPlan:
+    if execution_depth > 5:
         raise PipelineExecutionError("Subflows are limited to five nested levels.")
     validation = validate_graph(graph)
-    errors = [item for item in validation if item.severity == "error"]
-    if errors:
-        raise PipelineExecutionError(errors[0].message, node_id=errors[0].node_id)
-
+    first_error = next(
+        (item for item in validation if item.severity == "error"),
+        None,
+    )
+    if first_error is not None:
+        raise PipelineExecutionError(
+            first_error.message,
+            node_id=first_error.node_id,
+        )
     node_by_id = {node.id: node for node in graph.nodes}
     if preview_node_id is not None and preview_node_id not in node_by_id:
         raise PipelineExecutionError(
             "The requested preview node is not part of this pipeline.",
             node_id=preview_node_id,
         )
-    inputs = graph_inputs_by_port(graph)
     ordered, cyclic = topological_order(graph)
     if cyclic:
         raise PipelineExecutionError("Pipeline graph contains a cycle")
+    return _PreviewPlan(
+        graph=graph,
+        node_by_id=node_by_id,
+        inputs_by_port=graph_inputs_by_port(graph),
+        ordered_node_ids=ordered,
+        validation=validation,
+        preview_node_id=preview_node_id,
+        row_limit=row_limit,
+    )
 
-    outputs: dict[str, list[dict[str, Any]]] = {}
-    node_diagnostics: list[NodePreviewDiagnostic] = []
-    source_fingerprints: list[dict[str, Any]] = []
-    started = time.monotonic()
-    input_row_count = 0
 
-    for node_id in ordered:
-        node = node_by_id[node_id]
-        node_started = time.monotonic()
-        node_inputs = inputs.get(node_id, {})
-        input_sets = [
-            outputs[source_id]
-            for port_sources in node_inputs.values()
-            for source_id in port_sources
-        ]
-        input_rows = input_sets[0] if len(input_sets) == 1 else []
-        node_messages: list[str] = []
-        if time.monotonic() - started > MAX_EXECUTION_SECONDS:
-            failed = NodePreviewDiagnostic(
-                node_id=node.id,
-                status="failed",
-                input_rows=sum(len(rows) for rows in input_sets),
-                output_rows=0,
-                duration_ms=0,
-                columns=[],
-                messages=["Preview exceeded the two-second execution limit"],
-            )
-            raise PipelineExecutionError(
-                "Preview exceeded the two-second execution limit",
-                node_id=node.id,
-                node_diagnostics=tuple([*node_diagnostics, failed]),
-                source_fingerprints=tuple(source_fingerprints),
-                input_row_count=input_row_count,
-                diagnostics=tuple(item for item in validation if item.severity != "error"),
-                node_preview=_node_preview(outputs, preview_node_id, row_limit),
-            )
-        try:
-            executor = OPERATOR_REGISTRY.executor(node.type)
-            if executor is None:
-                raise PipelineExecutionError(
-                    f"Node type {node.type!r} has no registered executor.",
-                    node_id=node.id,
-                )
-            execution = executor(
-                OperatorExecutionContext(
-                    node=node,
-                    inputs_by_port=node_inputs,
-                    outputs=outputs,
-                    input_sets=input_sets,
-                    input_rows=input_rows,
-                    source_resolver=source_resolver,
-                    execution_depth=_execution_depth,
-                )
-            )
-            output_rows = execution.rows
-            input_row_count += execution.input_row_count
-            source_fingerprints.extend(execution.source_fingerprints)
-            node_messages.extend(execution.messages)
-            if node.type == "source.reference" and execution.messages:
-                validation.extend(
-                    DataflowDiagnostic(
-                        severity="warning",
-                        code="source.preview_truncated",
-                        message=message,
-                        node_id=node.id,
-                    )
-                    for message in execution.messages
-                )
-            if len(json.dumps(output_rows, default=str).encode("utf-8")) > MAX_RESULT_BYTES:
-                raise PipelineExecutionError(
-                    "A preview node exceeded the one-megabyte result limit.",
-                    node_id=node.id,
-                )
-        except (PipelineExecutionError, ArithmeticError, KeyError, TypeError, ValueError) as exc:
-            execution_error = (
-                exc
-                if isinstance(exc, PipelineExecutionError)
-                else PipelineExecutionError(str(exc), node_id=node.id)
-            )
-            failed = NodePreviewDiagnostic(
-                node_id=execution_error.node_id or node.id,
-                status="failed",
-                input_rows=sum(len(rows) for rows in input_sets),
-                output_rows=0,
-                duration_ms=round((time.monotonic() - node_started) * 1000, 3),
-                columns=[],
-                messages=[str(execution_error)],
-            )
-            raise PipelineExecutionError(
-                str(execution_error),
-                node_id=execution_error.node_id or node.id,
-                node_diagnostics=tuple([*node_diagnostics, failed]),
-                source_fingerprints=tuple(source_fingerprints),
-                input_row_count=input_row_count,
-                diagnostics=tuple(item for item in validation if item.severity != "error"),
-                node_preview=_node_preview(outputs, preview_node_id, row_limit),
-            ) from exc
+def _execute_preview_node(
+    plan: _PreviewPlan,
+    state: _PreviewState,
+    node: GraphNode,
+    *,
+    source_resolver: SourceResolver | None,
+    execution_depth: int,
+) -> None:
+    node_started = time.monotonic()
+    node_inputs = plan.inputs_by_port.get(node.id, {})
+    input_sets = [
+        state.outputs[source_id]
+        for port_sources in node_inputs.values()
+        for source_id in port_sources
+    ]
+    _enforce_preview_deadline(plan, state, node, input_sets)
+    try:
+        execution = _run_operator(
+            node,
+            node_inputs=node_inputs,
+            input_sets=input_sets,
+            outputs=state.outputs,
+            source_resolver=source_resolver,
+            execution_depth=execution_depth,
+        )
+        _record_execution(plan, state, node, execution)
+        _guard_result_size(execution.rows, node_id=node.id)
+    except (
+        PipelineExecutionError,
+        ArithmeticError,
+        KeyError,
+        TypeError,
+        ValueError,
+    ) as exc:
+        _raise_preview_node_failure(
+            plan,
+            state,
+            node,
+            input_sets=input_sets,
+            node_started=node_started,
+            exc=exc,
+        )
+    state.outputs[node.id] = execution.rows
+    state.node_diagnostics.append(
+        _successful_node_diagnostic(
+            node,
+            execution,
+            input_sets=input_sets,
+            node_started=node_started,
+        )
+    )
 
-        outputs[node.id] = output_rows
-        node_diagnostics.append(
-            NodePreviewDiagnostic(
+
+
+def _run_operator(
+    node: GraphNode,
+    *,
+    node_inputs: dict[str, list[str]],
+    input_sets: list[list[dict[str, Any]]],
+    outputs: dict[str, list[dict[str, Any]]],
+    source_resolver: SourceResolver | None,
+    execution_depth: int,
+) -> OperatorExecutionResult:
+    executor = OPERATOR_REGISTRY.executor(node.type)
+    if executor is None:
+        raise PipelineExecutionError(
+            f"Node type {node.type!r} has no registered executor.",
+            node_id=node.id,
+        )
+    return executor(
+        OperatorExecutionContext(
+            node=node,
+            inputs_by_port=node_inputs,
+            outputs=outputs,
+            input_sets=input_sets,
+            input_rows=input_sets[0] if len(input_sets) == 1 else [],
+            source_resolver=source_resolver,
+            execution_depth=execution_depth,
+        )
+    )
+
+
+def _record_execution(
+    plan: _PreviewPlan,
+    state: _PreviewState,
+    node: GraphNode,
+    execution: OperatorExecutionResult,
+) -> None:
+    state.input_row_count += execution.input_row_count
+    state.source_fingerprints.extend(execution.source_fingerprints)
+    if node.type == "source.reference":
+        plan.validation.extend(
+            DataflowDiagnostic(
+                severity="warning",
+                code="source.preview_truncated",
+                message=message,
                 node_id=node.id,
-                status="succeeded",
-                input_rows=sum(len(rows) for rows in input_sets),
-                output_rows=len(output_rows),
-                duration_ms=round((time.monotonic() - node_started) * 1000, 3),
-                columns=infer_columns(output_rows),
-                messages=node_messages,
             )
+            for message in execution.messages
         )
 
-    output_node = next(node for node in graph.nodes if node.type == "output")
-    all_rows = outputs[output_node.id]
-    rows = all_rows[:row_limit]
+
+def _guard_result_size(
+    rows: list[dict[str, Any]],
+    *,
+    node_id: str,
+) -> None:
+    if len(json.dumps(rows, default=str).encode("utf-8")) <= MAX_RESULT_BYTES:
+        return
+    raise PipelineExecutionError(
+        "A preview node exceeded the one-megabyte result limit.",
+        node_id=node_id,
+    )
+
+
+def _enforce_preview_deadline(
+    plan: _PreviewPlan,
+    state: _PreviewState,
+    node: GraphNode,
+    input_sets: list[list[dict[str, Any]]],
+) -> None:
+    if time.monotonic() - state.started <= MAX_EXECUTION_SECONDS:
+        return
+    message = "Preview exceeded the two-second execution limit"
+    failed = NodePreviewDiagnostic(
+        node_id=node.id,
+        status="failed",
+        input_rows=_input_row_total(input_sets),
+        output_rows=0,
+        duration_ms=0,
+        columns=[],
+        messages=[message],
+    )
+    raise PipelineExecutionError(
+        message,
+        node_id=node.id,
+        node_diagnostics=tuple([*state.node_diagnostics, failed]),
+        source_fingerprints=tuple(state.source_fingerprints),
+        input_row_count=state.input_row_count,
+        diagnostics=tuple(_warnings(plan.validation)),
+        node_preview=_node_preview(
+            state.outputs,
+            plan.preview_node_id,
+            plan.row_limit,
+        ),
+    )
+
+
+def _raise_preview_node_failure(
+    plan: _PreviewPlan,
+    state: _PreviewState,
+    node: GraphNode,
+    *,
+    input_sets: list[list[dict[str, Any]]],
+    node_started: float,
+    exc: BaseException,
+) -> NoReturn:
+    execution_error = (
+        exc
+        if isinstance(exc, PipelineExecutionError)
+        else PipelineExecutionError(str(exc), node_id=node.id)
+    )
+    failed = NodePreviewDiagnostic(
+        node_id=execution_error.node_id or node.id,
+        status="failed",
+        input_rows=_input_row_total(input_sets),
+        output_rows=0,
+        duration_ms=_elapsed_ms(node_started),
+        columns=[],
+        messages=[str(execution_error)],
+    )
+    raise PipelineExecutionError(
+        str(execution_error),
+        node_id=execution_error.node_id or node.id,
+        node_diagnostics=tuple([*state.node_diagnostics, failed]),
+        source_fingerprints=tuple(state.source_fingerprints),
+        input_row_count=state.input_row_count,
+        diagnostics=tuple(_warnings(plan.validation)),
+        node_preview=_node_preview(
+            state.outputs,
+            plan.preview_node_id,
+            plan.row_limit,
+        ),
+    ) from exc
+
+
+def _successful_node_diagnostic(
+    node: GraphNode,
+    execution: OperatorExecutionResult,
+    *,
+    input_sets: list[list[dict[str, Any]]],
+    node_started: float,
+) -> NodePreviewDiagnostic:
+    return NodePreviewDiagnostic(
+        node_id=node.id,
+        status="succeeded",
+        input_rows=_input_row_total(input_sets),
+        output_rows=len(execution.rows),
+        duration_ms=_elapsed_ms(node_started),
+        columns=infer_columns(execution.rows),
+        messages=list(execution.messages),
+    )
+
+
+def _preview_result(
+    plan: _PreviewPlan,
+    state: _PreviewState,
+) -> PipelineExecutionResult:
+    output_node = next(
+        node
+        for node in plan.graph.nodes
+        if node.type == "output"
+    )
+    all_rows = state.outputs[output_node.id]
+    rows = all_rows[: plan.row_limit]
     return PipelineExecutionResult(
         rows=rows,
         total_rows=len(all_rows),
         truncated=len(rows) < len(all_rows),
         columns=infer_columns(all_rows),
-        diagnostics=[item for item in validation if item.severity != "error"],
-        node_diagnostics=node_diagnostics,
-        node_preview=_node_preview(outputs, preview_node_id, row_limit),
-        source_fingerprints=source_fingerprints,
-        input_row_count=input_row_count,
+        diagnostics=_warnings(plan.validation),
+        node_diagnostics=state.node_diagnostics,
+        node_preview=_node_preview(
+            state.outputs,
+            plan.preview_node_id,
+            plan.row_limit,
+        ),
+        source_fingerprints=state.source_fingerprints,
+        input_row_count=state.input_row_count,
     )
 
 
+def _warnings(
+    diagnostics: list[DataflowDiagnostic],
+) -> list[DataflowDiagnostic]:
+    return [
+        item
+        for item in diagnostics
+        if item.severity != "error"
+    ]
+
+
+def _input_row_total(
+    input_sets: list[list[dict[str, Any]]],
+) -> int:
+    return sum(len(rows) for rows in input_sets)
+
+
+def _elapsed_ms(started: float) -> float:
+    return round((time.monotonic() - started) * 1000, 3)
+
+
 def _node_preview(
     outputs: dict[str, list[dict[str, Any]]],
     node_id: str | None,
@@ -397,32 +581,67 @@ def _derive_rows(
 
 
 def _derive_value(operation: str, values: list[Any], *, separator: str) -> Any:
-    if operation == "copy":
-        return values[0]
-    if operation == "upper":
-        return None if values[0] is None else str(values[0]).upper()
-    if operation == "lower":
-        return None if values[0] is None else str(values[0]).lower()
-    if operation == "trim":
-        return None if values[0] is None else str(values[0]).strip()
-    if operation == "concat":
-        return separator.join(str(value) for value in values if value is not None)
-    if operation == "coalesce":
-        return next((value for value in values if value not in (None, "")), None)
+    handler = _DERIVE_HANDLERS.get(operation)
+    if handler is None:
+        raise ValueError(f"unknown derive operation {operation!r}")
+    return handler(values, separator)
+
+
+def _unary_text(
+    values: list[Any],
+    _separator: str,
+    operation: Callable[[str], str],
+) -> str | None:
+    return None if values[0] is None else operation(str(values[0]))
+
+
+def _derive_concat(values: list[Any], separator: str) -> str:
+    return separator.join(str(value) for value in values if value is not None)
+
+
+def _derive_coalesce(values: list[Any], _separator: str) -> Any:
+    return next((value for value in values if value not in (None, "")), None)
+
+
+def _derive_numeric(
+    values: list[Any],
+    _separator: str,
+    operation: Callable[[Any, Any], Any],
+) -> Any:
     if any(value is None for value in values):
         return None
     left, right = values
     if isinstance(left, bool) or isinstance(right, bool):
         raise TypeError("boolean values are not numeric inputs")
-    if operation == "add":
-        return left + right
-    if operation == "subtract":
-        return left - right
-    if operation == "multiply":
-        return left * right
-    if operation == "divide":
-        return left / right
-    raise ValueError(f"unknown derive operation {operation!r}")
+    return operation(left, right)
+
+
+_DERIVE_HANDLERS: dict[str, Callable[[list[Any], str], Any]] = {
+    "copy": lambda values, _separator: values[0],
+    "upper": lambda values, separator: _unary_text(
+        values, separator, str.upper
+    ),
+    "lower": lambda values, separator: _unary_text(
+        values, separator, str.lower
+    ),
+    "trim": lambda values, separator: _unary_text(
+        values, separator, str.strip
+    ),
+    "concat": _derive_concat,
+    "coalesce": _derive_coalesce,
+    "add": lambda values, separator: _derive_numeric(
+        values, separator, lambda left, right: left + right
+    ),
+    "subtract": lambda values, separator: _derive_numeric(
+        values, separator, lambda left, right: left - right
+    ),
+    "multiply": lambda values, separator: _derive_numeric(
+        values, separator, lambda left, right: left * right
+    ),
+    "divide": lambda values, separator: _derive_numeric(
+        values, separator, lambda left, right: left / right
+    ),
+}
 
 
 def _expression_filter_rows(
@@ -848,46 +1067,95 @@ def _aggregate_rows(
 ) -> list[dict[str, Any]]:
     group_by = [str(item) for item in config.get("group_by", [])]
     aggregates = list(config["aggregates"])
+    grouped = _group_rows(rows, group_by=group_by)
+    return [
+        _aggregate_group(
+            key,
+            group_rows,
+            group_by=group_by,
+            aggregates=aggregates,
+            node_id=node_id,
+        )
+        for key, group_rows in grouped.items()
+    ]
+
+
+def _group_rows(
+    rows: list[dict[str, Any]],
+    *,
+    group_by: list[str],
+) -> dict[tuple[Any, ...], list[dict[str, Any]]]:
     grouped: dict[tuple[Any, ...], list[dict[str, Any]]] = defaultdict(list)
     if rows:
         for row in rows:
             grouped[tuple(row.get(column) for column in group_by)].append(row)
     elif not group_by:
         grouped[()] = []
+    return grouped
 
-    output: list[dict[str, Any]] = []
-    for key, group_rows in grouped.items():
-        result = {column: value for column, value in zip(group_by, key, strict=True)}
-        for aggregate in aggregates:
-            function = str(aggregate["function"])
-            column = aggregate.get("column")
-            alias = str(aggregate["alias"])
-            values = [row.get(column) for row in group_rows if column is not None and row.get(column) is not None]
-            try:
-                if function == "count":
-                    result[alias] = len(group_rows) if column in (None, "", "*") else len(values)
-                elif function == "sum":
-                    result[alias] = sum(values) if values else 0
-                elif function == "avg":
-                    result[alias] = sum(values) / len(values) if values else None
-                elif function == "min":
-                    result[alias] = min(values) if values else None
-                elif function == "max":
-                    result[alias] = max(values) if values else None
-            except (ArithmeticError, TypeError, ValueError) as exc:
-                raise PipelineExecutionError(
-                    f"Cannot calculate {function.upper()} for {column!r}: {exc}",
-                    node_id=node_id,
-                ) from exc
-        output.append(result)
-    return output
+
+def _aggregate_group(
+    key: tuple[Any, ...],
+    rows: list[dict[str, Any]],
+    *,
+    group_by: list[str],
+    aggregates: list[dict[str, Any]],
+    node_id: str,
+) -> dict[str, Any]:
+    result = {
+        column: value
+        for column, value in zip(group_by, key, strict=True)
+    }
+    for aggregate in aggregates:
+        alias = str(aggregate["alias"])
+        result[alias] = _aggregate_value(
+            str(aggregate["function"]),
+            aggregate.get("column"),
+            rows,
+            node_id=node_id,
+        )
+    return result
+
+
+def _aggregate_value(
+    function: str,
+    column: Any,
+    rows: list[dict[str, Any]],
+    *,
+    node_id: str,
+) -> Any:
+    values = [
+        row.get(column)
+        for row in rows
+        if column is not None and row.get(column) is not None
+    ]
+    try:
+        if function == "count":
+            return len(rows) if column in (None, "", "*") else len(values)
+        handler = _AGGREGATE_HANDLERS.get(function)
+        if handler is None:
+            raise ValueError(f"unknown aggregate function {function!r}")
+        return handler(values)
+    except (ArithmeticError, TypeError, ValueError) as exc:
+        raise PipelineExecutionError(
+            f"Cannot calculate {function.upper()} for {column!r}: {exc}",
+            node_id=node_id,
+        ) from exc
+
+
+_AGGREGATE_HANDLERS: dict[str, Callable[[list[Any]], Any]] = {
+    "sum": lambda values: sum(values) if values else 0,
+    "avg": lambda values: sum(values) / len(values) if values else None,
+    "min": lambda values: min(values) if values else None,
+    "max": lambda values: max(values) if values else None,
+}
 
 
 def _sort_rows(rows: list[dict[str, Any]], config: dict[str, Any]) -> list[dict[str, Any]]:
     result = [dict(row) for row in rows]
-    for field in reversed(config["fields"]):
-        column = str(field["column"])
-        reverse = field.get("direction", "asc") == "desc"
+    for field_config in reversed(config["fields"]):
+        column = str(field_config["column"])
+        reverse = field_config.get("direction", "asc") == "desc"
         concrete = [row for row in result if row.get(column) is not None]
         nulls = [row for row in result if row.get(column) is None]
         concrete.sort(key=lambda row: _sortable_value(row[column]), reverse=reverse)

src/govoplan_dataflow/backend/expressions.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
+import operator
 from dataclasses import dataclass
 from decimal import Decimal
-from typing import Any, Literal
+from typing import Any, Callable, Literal
 
 import sqlglot
 from sqlglot import exp
@@ -141,133 +142,255 @@ def infer_expression_type(
 
 
 def _evaluate(expression: exp.Expression, row: dict[str, Any]) -> Any:
-    if isinstance(expression, exp.Paren):
-        return _evaluate(expression.this, row)
-    if isinstance(expression, exp.Column):
-        return row.get(expression.name)
-    if isinstance(expression, exp.Null):
-        return None
-    if isinstance(expression, exp.Boolean):
-        return bool(expression.this)
-    if isinstance(expression, exp.Literal):
-        return _literal(expression)
-    if isinstance(expression, exp.Neg):
-        value = _evaluate(expression.this, row)
-        return None if value is None else -value
-    if isinstance(expression, exp.Not):
-        return not bool(_evaluate(expression.this, row))
+    evaluator = _EVALUATORS.get(type(expression))
+    if evaluator is None:
+        raise ExpressionError(
+            f"Expression element {type(expression).__name__} cannot be evaluated."
+        )
+    return evaluator(expression, row)
+
+
+def _evaluate_child(expression: exp.Expression, row: dict[str, Any]) -> Any:
+    return _evaluate(expression.this, row)
+
+
+def _evaluate_negation(expression: exp.Expression, row: dict[str, Any]) -> Any:
+    value = _evaluate_child(expression, row)
+    return None if value is None else -value
+
+
+def _evaluate_boolean_not(
+    expression: exp.Expression,
+    row: dict[str, Any],
+) -> bool:
+    return not bool(_evaluate_child(expression, row))
+
+
+def _evaluate_boolean_binary(
+    expression: exp.Expression,
+    row: dict[str, Any],
+) -> bool:
+    left = bool(_evaluate(expression.this, row))
     if isinstance(expression, exp.And):
-        return bool(_evaluate(expression.this, row)) and bool(
-            _evaluate(expression.expression, row)
-        )
-    if isinstance(expression, exp.Or):
-        return bool(_evaluate(expression.this, row)) or bool(
-            _evaluate(expression.expression, row)
-        )
-    if isinstance(expression, exp.Is):
-        left = _evaluate(expression.this, row)
-        right = _evaluate(expression.expression, row)
-        return left is right if right is None else left == right
-    if isinstance(expression, (exp.Add, exp.Sub, exp.Mul, exp.Div, exp.Mod)):
-        left = _evaluate(expression.this, row)
-        right = _evaluate(expression.expression, row)
-        if left is None or right is None:
-            return None
-        if isinstance(expression, exp.Add):
-            return left + right
-        if isinstance(expression, exp.Sub):
-            return left - right
-        if isinstance(expression, exp.Mul):
-            return left * right
-        if isinstance(expression, exp.Div):
-            return left / right
-        return left % right
-    if isinstance(expression, (exp.EQ, exp.NEQ, exp.GT, exp.GTE, exp.LT, exp.LTE)):
-        left = _evaluate(expression.this, row)
-        right = _evaluate(expression.expression, row)
-        if isinstance(expression, exp.EQ):
-            return left == right
-        if isinstance(expression, exp.NEQ):
-            return left != right
-        if left is None or right is None:
-            return False
-        if isinstance(expression, exp.GT):
-            return left > right
-        if isinstance(expression, exp.GTE):
-            return left >= right
-        if isinstance(expression, exp.LT):
-            return left < right
-        return left <= right
-    if isinstance(expression, exp.Lower):
-        return _string_call(expression.this, row, str.lower)
-    if isinstance(expression, exp.Upper):
-        return _string_call(expression.this, row, str.upper)
-    if isinstance(expression, exp.Trim):
-        return _string_call(expression.this, row, str.strip)
-    if isinstance(expression, exp.Length):
-        value = _evaluate(expression.this, row)
-        return None if value is None else len(value)
-    if isinstance(expression, exp.Abs):
-        value = _evaluate(expression.this, row)
-        return None if value is None else abs(value)
-    if isinstance(expression, exp.Round):
-        value = _evaluate(expression.this, row)
-        decimals = expression.args.get("decimals")
-        places = int(_evaluate(decimals, row)) if decimals is not None else 0
-        return None if value is None else round(value, places)
-    if isinstance(expression, exp.Coalesce):
-        arguments = [expression.this, *expression.expressions]
-        return next(
-            (
-                value
-                for argument in arguments
-                if (value := _evaluate(argument, row)) is not None
-            ),
-            None,
-        )
-    if isinstance(expression, exp.Concat):
-        return "".join(
-            "" if (value := _evaluate(item, row)) is None else str(value)
-            for item in expression.expressions
-        )
-    if isinstance(expression, exp.Replace):
-        value = _evaluate(expression.this, row)
-        old = _evaluate(expression.expression, row)
-        new = _evaluate(expression.args["replacement"], row)
-        return None if value is None else str(value).replace(str(old), str(new))
-    if isinstance(expression, exp.Substring):
-        value = _evaluate(expression.this, row)
-        if value is None:
-            return None
-        start = int(_evaluate(expression.args["start"], row)) - 1
-        length = expression.args.get("length")
-        return (
-            str(value)[start:]
-            if length is None
-            else str(value)[start : start + int(_evaluate(length, row))]
-        )
-    if isinstance(expression, exp.Cast):
-        return convert_value(
-            _evaluate(expression.this, row),
-            _cast_type(expression),
-            on_error="fail",
-        )
-    if isinstance(expression, exp.Case):
-        base = _evaluate(expression.this, row) if expression.this is not None else None
-        for condition in expression.args.get("ifs") or []:
-            if not isinstance(condition, exp.If):
-                continue
-            candidate = _evaluate(condition.this, row)
-            matched = candidate == base if expression.this is not None else bool(candidate)
-            if matched:
-                return _evaluate(condition.args["true"], row)
-        default = expression.args.get("default")
-        return _evaluate(default, row) if default is not None else None
-    raise ExpressionError(
-        f"Expression element {type(expression).__name__} cannot be evaluated."
+        return left and bool(_evaluate(expression.expression, row))
+    return left or bool(_evaluate(expression.expression, row))
+
+
+def _evaluate_is(expression: exp.Expression, row: dict[str, Any]) -> bool:
+    left = _evaluate(expression.this, row)
+    right = _evaluate(expression.expression, row)
+    return left is right if right is None else left == right
+
+
+def _evaluate_binary(
+    expression: exp.Expression,
+    row: dict[str, Any],
+    operations: dict[type[exp.Expression], Callable[[Any, Any], Any]],
+    *,
+    null_result: Any,
+) -> Any:
+    left = _evaluate(expression.this, row)
+    right = _evaluate(expression.expression, row)
+    if left is None or right is None:
+        return null_result
+    return operations[type(expression)](left, right)
+
+
+def _evaluate_arithmetic(expression: exp.Expression, row: dict[str, Any]) -> Any:
+    return _evaluate_binary(
+        expression,
+        row,
+        _ARITHMETIC_OPERATIONS,
+        null_result=None,
     )
 
 
+def _evaluate_comparison(
+    expression: exp.Expression,
+    row: dict[str, Any],
+) -> bool:
+    if isinstance(expression, (exp.EQ, exp.NEQ)):
+        left = _evaluate(expression.this, row)
+        right = _evaluate(expression.expression, row)
+        return _COMPARISON_OPERATIONS[type(expression)](left, right)
+    return _evaluate_binary(
+        expression,
+        row,
+        _COMPARISON_OPERATIONS,
+        null_result=False,
+    )
+
+
+def _evaluate_string_function(
+    expression: exp.Expression,
+    row: dict[str, Any],
+) -> str | None:
+    return _string_call(
+        expression.this,
+        row,
+        _STRING_OPERATIONS[type(expression)],
+    )
+
+
+def _evaluate_length(expression: exp.Expression, row: dict[str, Any]) -> int | None:
+    value = _evaluate_child(expression, row)
+    return None if value is None else len(value)
+
+
+def _evaluate_abs(expression: exp.Expression, row: dict[str, Any]) -> Any:
+    value = _evaluate_child(expression, row)
+    return None if value is None else abs(value)
+
+
+def _evaluate_round(expression: exp.Expression, row: dict[str, Any]) -> Any:
+    value = _evaluate_child(expression, row)
+    decimals = expression.args.get("decimals")
+    places = int(_evaluate(decimals, row)) if decimals is not None else 0
+    return None if value is None else round(value, places)
+
+
+def _evaluate_coalesce(expression: exp.Expression, row: dict[str, Any]) -> Any:
+    arguments = [expression.this, *expression.expressions]
+    return next(
+        (
+            value
+            for argument in arguments
+            if (value := _evaluate(argument, row)) is not None
+        ),
+        None,
+    )
+
+
+def _evaluate_concat(expression: exp.Expression, row: dict[str, Any]) -> str:
+    return "".join(
+        "" if (value := _evaluate(item, row)) is None else str(value)
+        for item in expression.expressions
+    )
+
+
+def _evaluate_replace(expression: exp.Expression, row: dict[str, Any]) -> Any:
+    value = _evaluate_child(expression, row)
+    old = _evaluate(expression.expression, row)
+    new = _evaluate(expression.args["replacement"], row)
+    return None if value is None else str(value).replace(str(old), str(new))
+
+
+def _evaluate_substring(expression: exp.Expression, row: dict[str, Any]) -> Any:
+    value = _evaluate_child(expression, row)
+    if value is None:
+        return None
+    start = int(_evaluate(expression.args["start"], row)) - 1
+    length = expression.args.get("length")
+    if length is None:
+        return str(value)[start:]
+    return str(value)[start : start + int(_evaluate(length, row))]
+
+
+def _evaluate_cast(expression: exp.Expression, row: dict[str, Any]) -> Any:
+    return convert_value(
+        _evaluate_child(expression, row),
+        _cast_type(expression),  # type: ignore[arg-type]
+        on_error="fail",
+    )
+
+
+def _evaluate_case(expression: exp.Expression, row: dict[str, Any]) -> Any:
+    base_expression = expression.this
+    base = _evaluate(base_expression, row) if base_expression is not None else None
+    for condition in expression.args.get("ifs") or []:
+        if isinstance(condition, exp.If) and _case_matches(
+            condition,
+            row,
+            base=base,
+            has_base=base_expression is not None,
+        ):
+            return _evaluate(condition.args["true"], row)
+    default = expression.args.get("default")
+    return _evaluate(default, row) if default is not None else None
+
+
+def _case_matches(
+    condition: exp.If,
+    row: dict[str, Any],
+    *,
+    base: Any,
+    has_base: bool,
+) -> bool:
+    candidate = _evaluate(condition.this, row)
+    return candidate == base if has_base else bool(candidate)
+
+
+def _evaluate_if(expression: exp.Expression, row: dict[str, Any]) -> Any:
+    branch = "true" if bool(_evaluate(expression.this, row)) else "false"
+    selected = expression.args.get(branch)
+    return _evaluate(selected, row) if selected is not None else None
+
+
+_ARITHMETIC_OPERATIONS: dict[
+    type[exp.Expression],
+    Callable[[Any, Any], Any],
+] = {
+    exp.Add: operator.add,
+    exp.Sub: operator.sub,
+    exp.Mul: operator.mul,
+    exp.Div: operator.truediv,
+    exp.Mod: operator.mod,
+}
+_COMPARISON_OPERATIONS: dict[
+    type[exp.Expression],
+    Callable[[Any, Any], bool],
+] = {
+    exp.EQ: operator.eq,
+    exp.NEQ: operator.ne,
+    exp.GT: operator.gt,
+    exp.GTE: operator.ge,
+    exp.LT: operator.lt,
+    exp.LTE: operator.le,
+}
+_STRING_OPERATIONS: dict[type[exp.Expression], Callable[[str], str]] = {
+    exp.Lower: str.lower,
+    exp.Upper: str.upper,
+    exp.Trim: str.strip,
+}
+_EVALUATORS: dict[
+    type[exp.Expression],
+    Callable[[exp.Expression, dict[str, Any]], Any],
+] = {
+    exp.Paren: _evaluate_child,
+    exp.Column: lambda expression, row: row.get(expression.name),
+    exp.Null: lambda _expression, _row: None,
+    exp.Boolean: lambda expression, _row: bool(expression.this),
+    exp.Literal: lambda expression, _row: _literal(expression),  # type: ignore[arg-type]
+    exp.Neg: _evaluate_negation,
+    exp.Not: _evaluate_boolean_not,
+    exp.And: _evaluate_boolean_binary,
+    exp.Or: _evaluate_boolean_binary,
+    exp.Is: _evaluate_is,
+    **{
+        expression_type: _evaluate_arithmetic
+        for expression_type in _ARITHMETIC_OPERATIONS
+    },
+    **{
+        expression_type: _evaluate_comparison
+        for expression_type in _COMPARISON_OPERATIONS
+    },
+    **{
+        expression_type: _evaluate_string_function
+        for expression_type in _STRING_OPERATIONS
+    },
+    exp.Length: _evaluate_length,
+    exp.Abs: _evaluate_abs,
+    exp.Round: _evaluate_round,
+    exp.Coalesce: _evaluate_coalesce,
+    exp.Concat: _evaluate_concat,
+    exp.Replace: _evaluate_replace,
+    exp.Substring: _evaluate_substring,
+    exp.Cast: _evaluate_cast,
+    exp.Case: _evaluate_case,
+    exp.If: _evaluate_if,
+}
+
+
 def convert_value(
     value: Any,
     target_type: str,
@@ -312,19 +435,115 @@ def _infer(
     expression: exp.Expression,
     schema: dict[str, ExpressionDataType],
 ) -> ExpressionDataType:
-    if isinstance(expression, exp.Column):
-        return schema.get(expression.name, "unknown")
-    if isinstance(expression, exp.Null):
-        return "null"
-    if isinstance(expression, exp.Boolean):
-        return "boolean"
-    if isinstance(expression, exp.Literal):
-        if expression.is_string:
-            return "string"
-        return "number" if "." in str(expression.this) else "integer"
-    if isinstance(
-        expression,
-        (
+    inferer = _TYPE_INFERERS.get(type(expression))
+    return inferer(expression, schema) if inferer is not None else "unknown"
+
+
+def _infer_literal(
+    expression: exp.Expression,
+    _schema: dict[str, ExpressionDataType],
+) -> ExpressionDataType:
+    literal = expression
+    if not isinstance(literal, exp.Literal):
+        return "unknown"
+    if literal.is_string:
+        return "string"
+    return "number" if "." in str(literal.this) else "integer"
+
+
+def _infer_cast(
+    expression: exp.Expression,
+    _schema: dict[str, ExpressionDataType],
+) -> ExpressionDataType:
+    return _cast_type(expression)  # type: ignore[arg-type,return-value]
+
+
+def _infer_numeric(
+    expression: exp.Expression,
+    schema: dict[str, ExpressionDataType],
+) -> ExpressionDataType:
+    child_types = {
+        _infer(item, schema)
+        for item in expression.iter_expressions()
+    }
+    if "number" in child_types or isinstance(expression, exp.Div):
+        return "number"
+    return "integer"
+
+
+def _infer_coalesce(
+    expression: exp.Expression,
+    schema: dict[str, ExpressionDataType],
+) -> ExpressionDataType:
+    candidates = (
+        _infer(item, schema)
+        for item in (expression.this, *expression.expressions)
+    )
+    return next(
+        (item for item in candidates if item not in {"null", "unknown"}),
+        "unknown",
+    )
+
+
+def _infer_case(
+    expression: exp.Expression,
+    schema: dict[str, ExpressionDataType],
+) -> ExpressionDataType:
+    types = [
+        _infer(item.args["true"], schema)
+        for item in expression.args.get("ifs") or []
+        if isinstance(item, exp.If)
+    ]
+    default = expression.args.get("default")
+    if default is not None:
+        types.append(_infer(default, schema))
+    return _common_expression_type(types)
+
+
+def _infer_if(
+    expression: exp.Expression,
+    schema: dict[str, ExpressionDataType],
+) -> ExpressionDataType:
+    return _common_expression_type(
+        [
+            _infer(branch, schema)
+            for key in ("true", "false")
+            if (branch := expression.args.get(key)) is not None
+        ]
+    )
+
+
+def _common_expression_type(
+    types: list[ExpressionDataType],
+) -> ExpressionDataType:
+    concrete = {item for item in types if item not in {"null", "unknown"}}
+    return concrete.pop() if len(concrete) == 1 else "unknown"
+
+
+def _infer_child(
+    expression: exp.Expression,
+    schema: dict[str, ExpressionDataType],
+) -> ExpressionDataType:
+    return _infer(expression.this, schema)
+
+
+_TYPE_INFERERS: dict[
+    type[exp.Expression],
+    Callable[
+        [exp.Expression, dict[str, ExpressionDataType]],
+        ExpressionDataType,
+    ],
+] = {
+    exp.Column: lambda expression, schema: schema.get(
+        expression.name,
+        "unknown",
+    ),
+    exp.Null: lambda _expression, _schema: "null",
+    exp.Boolean: lambda _expression, _schema: "boolean",
+    exp.Literal: _infer_literal,
+    **{
+        expression_type: lambda _expression, _schema: "boolean"
+        for expression_type in (
             exp.EQ,
             exp.NEQ,
             exp.GT,
@@ -335,41 +554,39 @@ def _infer(
             exp.Or,
             exp.Is,
             exp.Not,
-        ),
-    ):
-        return "boolean"
-    if isinstance(expression, (exp.Length,)):
-        return "integer"
-    if isinstance(expression, (exp.Lower, exp.Upper, exp.Trim, exp.Concat, exp.Replace, exp.Substring)):
-        return "string"
-    if isinstance(expression, exp.Cast):
-        return _cast_type(expression)  # type: ignore[return-value]
-    if isinstance(expression, (exp.Add, exp.Sub, exp.Mul, exp.Div, exp.Mod, exp.Abs, exp.Round)):
-        child_types = {
-            _infer(item, schema)
-            for item in expression.iter_expressions()
-        }
-        return "number" if "number" in child_types or isinstance(expression, exp.Div) else "integer"
-    if isinstance(expression, exp.Coalesce):
-        types = [
-            _infer(item, schema)
-            for item in (expression.this, *expression.expressions)
-        ]
-        return next((item for item in types if item not in {"null", "unknown"}), "unknown")
-    if isinstance(expression, exp.Case):
-        types = [
-            _infer(item.args["true"], schema)
-            for item in expression.args.get("ifs") or []
-            if isinstance(item, exp.If)
-        ]
-        default = expression.args.get("default")
-        if default is not None:
-            types.append(_infer(default, schema))
-        concrete = {item for item in types if item not in {"null", "unknown"}}
-        return concrete.pop() if len(concrete) == 1 else "unknown"
-    if isinstance(expression, (exp.Paren, exp.Neg)):
-        return _infer(expression.this, schema)
-    return "unknown"
+        )
+    },
+    exp.Length: lambda _expression, _schema: "integer",
+    **{
+        expression_type: lambda _expression, _schema: "string"
+        for expression_type in (
+            exp.Lower,
+            exp.Upper,
+            exp.Trim,
+            exp.Concat,
+            exp.Replace,
+            exp.Substring,
+        )
+    },
+    exp.Cast: _infer_cast,
+    **{
+        expression_type: _infer_numeric
+        for expression_type in (
+            exp.Add,
+            exp.Sub,
+            exp.Mul,
+            exp.Div,
+            exp.Mod,
+            exp.Abs,
+            exp.Round,
+        )
+    },
+    exp.Coalesce: _infer_coalesce,
+    exp.Case: _infer_case,
+    exp.If: _infer_if,
+    exp.Paren: _infer_child,
+    exp.Neg: _infer_child,
+}
 
 
 def _literal(expression: exp.Literal) -> Any:

src/govoplan_dataflow/backend/schema_validation.py

@@ -0,0 +1,892 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from govoplan_dataflow.backend.expressions import (
+    ExpressionError,
+    infer_expression_type,
+    parse_expression,
+)
+from govoplan_dataflow.backend.operator_registry import OPERATOR_REGISTRY
+from govoplan_dataflow.backend.schemas import (
+    DataflowDiagnostic,
+    GraphNode,
+    PipelineGraph,
+)
+
+
+@dataclass(frozen=True, slots=True)
+class SchemaState:
+    columns: frozenset[str]
+    open: bool = False
+    types: dict[str, str] = field(default_factory=dict)
+
+    def knows(self, column: str) -> bool:
+        return self.open or column in self.columns
+
+    def type_of(self, column: str) -> str:
+        return self.types.get(column, "unknown")
+
+
+@dataclass(frozen=True, slots=True)
+class SchemaPropagationContext:
+    node: GraphNode
+    input_state: SchemaState
+    input_states: tuple[SchemaState, ...]
+    inputs_by_port: dict[str, list[str]]
+    schemas: dict[str, SchemaState]
+
+    def port_state(self, port: str) -> SchemaState:
+        source_ids = self.inputs_by_port.get(port, ())
+        if not source_ids:
+            return unknown_schema()
+        return self.schemas.get(source_ids[0], unknown_schema())
+
+
+@dataclass(frozen=True, slots=True)
+class SchemaPropagationResult:
+    state: SchemaState
+    diagnostics: tuple[DataflowDiagnostic, ...] = ()
+
+
+def validate_graph_schemas(
+    graph: PipelineGraph,
+    *,
+    ordered: list[str],
+) -> list[DataflowDiagnostic]:
+    node_by_id = {node.id: node for node in graph.nodes}
+    inputs = _graph_inputs_by_port(graph)
+    schemas: dict[str, SchemaState] = {}
+    diagnostics: list[DataflowDiagnostic] = []
+    for node_id in ordered:
+        node = node_by_id[node_id]
+        context = _propagation_context(node, inputs, schemas)
+        propagator = OPERATOR_REGISTRY.schema_propagator(node.type)
+        if propagator is None:
+            diagnostics.append(
+                _error(
+                    "node.schema_propagator_missing",
+                    f"Node type {node.type!r} has no schema propagator.",
+                    node_id=node.id,
+                )
+            )
+            schemas[node.id] = context.input_state
+            continue
+        result = propagator(context)
+        if not isinstance(result, SchemaPropagationResult):
+            raise TypeError(
+                f"Schema propagator for {node.type!r} returned "
+                f"{type(result).__name__}, not SchemaPropagationResult."
+            )
+        schemas[node.id] = result.state
+        diagnostics.extend(result.diagnostics)
+    return diagnostics
+
+
+def _graph_inputs_by_port(
+    graph: PipelineGraph,
+) -> dict[str, dict[str, list[str]]]:
+    result: dict[str, dict[str, list[str]]] = {}
+    for edge in graph.edges:
+        result.setdefault(edge.target, {}).setdefault(
+            edge.target_port,
+            [],
+        ).append(edge.source)
+    return result
+
+
+def _propagation_context(
+    node: GraphNode,
+    inputs: dict[str, dict[str, list[str]]],
+    schemas: dict[str, SchemaState],
+) -> SchemaPropagationContext:
+    node_inputs = inputs.get(node.id, {})
+    input_states = tuple(
+        schemas[source_id]
+        for port_sources in node_inputs.values()
+        for source_id in port_sources
+        if source_id in schemas
+    )
+    return SchemaPropagationContext(
+        node=node,
+        input_state=input_states[0] if input_states else unknown_schema(),
+        input_states=input_states,
+        inputs_by_port=node_inputs,
+        schemas=schemas,
+    )
+
+
+def _inline_source(
+    context: SchemaPropagationContext,
+) -> SchemaPropagationResult:
+    return SchemaPropagationResult(_inline_schema(context.node.config.get("rows")))
+
+
+def _reference_source(
+    context: SchemaPropagationContext,
+) -> SchemaPropagationResult:
+    return SchemaPropagationResult(
+        _configured_schema(context.node.config.get("source_columns"))
+    )
+
+
+def _union(context: SchemaPropagationContext) -> SchemaPropagationResult:
+    diagnostics: list[DataflowDiagnostic] = []
+    closed_shapes = {
+        state.columns
+        for state in context.input_states
+        if not state.open
+    }
+    if len(closed_shapes) > 1:
+        diagnostics.append(
+            _warning(
+                "union.schema_mismatch",
+                "Appended inputs use different columns; missing values will be null.",
+                node_id=context.node.id,
+            )
+        )
+    return SchemaPropagationResult(
+        SchemaState(
+            frozenset().union(
+                *(state.columns for state in context.input_states)
+            ),
+            open=any(state.open for state in context.input_states),
+            types=_merged_schema_types(context.input_states),
+        ),
+        tuple(diagnostics),
+    )
+
+
+def _join(context: SchemaPropagationContext) -> SchemaPropagationResult:
+    node = context.node
+    left_state = context.port_state("left")
+    right_state = context.port_state("right")
+    diagnostics = [
+        *_unknown_columns(
+            node,
+            left_state,
+            node.config.get("left_keys"),
+            field_name="left_keys",
+        ),
+        *_unknown_columns(
+            node,
+            right_state,
+            node.config.get("right_keys"),
+            field_name="right_keys",
+        ),
+    ]
+    prefix = str(node.config.get("right_prefix", "right_"))
+    prefixed_right = {
+        f"{prefix}{column}"
+        for column in right_state.columns
+    }
+    collisions = left_state.columns & prefixed_right
+    if collisions:
+        diagnostics.append(
+            _error(
+                "join.output_collision",
+                f"Join output columns collide: {', '.join(sorted(collisions))}.",
+                node_id=node.id,
+                field="right_prefix",
+            )
+        )
+    return SchemaPropagationResult(
+        SchemaState(
+            frozenset(left_state.columns | prefixed_right),
+            open=left_state.open or right_state.open,
+            types={
+                **left_state.types,
+                **{
+                    f"{prefix}{column}": right_state.type_of(column)
+                    for column in right_state.columns
+                },
+            },
+        ),
+        tuple(diagnostics),
+    )
+
+
+def _filter(context: SchemaPropagationContext) -> SchemaPropagationResult:
+    return _passthrough_with_columns(
+        context,
+        [context.node.config.get("column")],
+        field_name="column",
+    )
+
+
+def _filter_expression(
+    context: SchemaPropagationContext,
+) -> SchemaPropagationResult:
+    parsed, diagnostics = _node_expression(context.node, "expression")
+    if parsed is not None:
+        diagnostics.extend(
+            _unknown_columns(
+                context.node,
+                context.input_state,
+                list(parsed.columns),
+                field_name="expression",
+            )
+        )
+    return SchemaPropagationResult(
+        context.input_state,
+        tuple(diagnostics),
+    )
+
+
+def _distinct(context: SchemaPropagationContext) -> SchemaPropagationResult:
+    return _passthrough_with_columns(
+        context,
+        context.node.config.get("columns"),
+        field_name="columns",
+    )
+
+
+def _select(context: SchemaPropagationContext) -> SchemaPropagationResult:
+    selected, output = _selected_columns(context.node.config.get("fields"))
+    diagnostics = [
+        *_unknown_columns(
+            context.node,
+            context.input_state,
+            selected,
+            field_name="fields",
+        ),
+        *_duplicate_outputs(
+            context.node,
+            output,
+            field_name="fields",
+        ),
+    ]
+    return SchemaPropagationResult(
+        SchemaState(
+            frozenset(output),
+            types={
+                target: context.input_state.type_of(source)
+                for source, target in zip(selected, output, strict=False)
+            },
+        ),
+        tuple(diagnostics),
+    )
+
+
+def _selected_columns(value: object) -> tuple[list[str], list[str]]:
+    selected: list[str] = []
+    output: list[str] = []
+    if not isinstance(value, list):
+        return selected, output
+    for item in value:
+        if isinstance(item, str):
+            selected.append(item)
+            output.append(item)
+        elif isinstance(item, dict):
+            source = item.get("column")
+            target = item.get("alias") or source
+            if isinstance(source, str):
+                selected.append(source)
+            if isinstance(target, str):
+                output.append(target)
+    return selected, output
+
+
+def _derive(context: SchemaPropagationContext) -> SchemaPropagationResult:
+    node = context.node
+    source_columns = _text_items(node.config.get("source_columns"))
+    diagnostics = _unknown_columns(
+        node,
+        context.input_state,
+        source_columns,
+        field_name="source_columns",
+    )
+    target = node.config.get("target_column")
+    if not isinstance(target, str) or not target:
+        return SchemaPropagationResult(
+            context.input_state,
+            tuple(diagnostics),
+        )
+    if target in context.input_state.columns:
+        diagnostics.append(
+            _warning(
+                "derive.overwrites_column",
+                f"Derived column {target!r} replaces an existing value.",
+                node_id=node.id,
+                field="target_column",
+            )
+        )
+    result_type = _derive_result_type(
+        str(node.config.get("operation") or ""),
+        [
+            context.input_state.type_of(column)
+            for column in source_columns
+        ],
+    )
+    return SchemaPropagationResult(
+        _with_column(context.input_state, target, result_type),
+        tuple(diagnostics),
+    )
+
+
+def _expression(context: SchemaPropagationContext) -> SchemaPropagationResult:
+    node = context.node
+    parsed, diagnostics = _node_expression(node, "expression")
+    inferred = "unknown"
+    if parsed is not None:
+        diagnostics.extend(
+            _unknown_columns(
+                node,
+                context.input_state,
+                list(parsed.columns),
+                field_name="expression",
+            )
+        )
+        inferred = infer_expression_type(
+            parsed,
+            {
+                name: context.input_state.type_of(name)  # type: ignore[dict-item]
+                for name in context.input_state.columns
+            },
+        )
+    expected = str(node.config.get("result_type") or "unknown")
+    if expected != "unknown" and inferred not in {"unknown", "null", expected}:
+        diagnostics.append(
+            _warning(
+                "expression.type_mismatch",
+                f"Expression infers {inferred}, not {expected}.",
+                node_id=node.id,
+                field="result_type",
+            )
+        )
+    target = str(node.config.get("target_column") or "")
+    state = (
+        _with_column(
+            context.input_state,
+            target,
+            expected if expected != "unknown" else inferred,
+        )
+        if target
+        else context.input_state
+    )
+    return SchemaPropagationResult(state, tuple(diagnostics))
+
+
+def _convert_or_replace(
+    context: SchemaPropagationContext,
+) -> SchemaPropagationResult:
+    node = context.node
+    source = str(node.config.get("source_column") or "")
+    target = str(node.config.get("target_column") or "")
+    diagnostics = _unknown_columns(
+        node,
+        context.input_state,
+        [source],
+        field_name="source_column",
+    )
+    target_type = (
+        str(node.config.get("target_type") or "unknown")
+        if node.type == "convert"
+        else context.input_state.type_of(source)
+    )
+    state = (
+        _with_column(context.input_state, target, target_type)
+        if target
+        else context.input_state
+    )
+    return SchemaPropagationResult(state, tuple(diagnostics))
+
+
+def _aggregate(context: SchemaPropagationContext) -> SchemaPropagationResult:
+    node = context.node
+    group_by = _text_items(node.config.get("group_by"))
+    aggregates = _mapping_items(node.config.get("aggregates"))
+    source_columns = [
+        str(item.get("column"))
+        for item in aggregates
+        if item.get("column") not in (None, "", "*")
+    ]
+    aliases = [
+        str(item.get("alias"))
+        for item in aggregates
+        if item.get("alias")
+    ]
+    output = [*group_by, *aliases]
+    diagnostics = [
+        *_unknown_columns(
+            node,
+            context.input_state,
+            [*group_by, *source_columns],
+            field_name="aggregates",
+        ),
+        *_duplicate_outputs(node, output, field_name="aggregates"),
+    ]
+    aggregate_types = {
+        str(item["alias"]): (
+            "integer"
+            if item.get("function") == "count"
+            else context.input_state.type_of(str(item.get("column") or ""))
+        )
+        for item in aggregates
+        if item.get("alias")
+    }
+    return SchemaPropagationResult(
+        SchemaState(
+            frozenset(output),
+            types={
+                **{
+                    column: context.input_state.type_of(column)
+                    for column in group_by
+                },
+                **aggregate_types,
+            },
+        ),
+        tuple(diagnostics),
+    )
+
+
+def _sort(context: SchemaPropagationContext) -> SchemaPropagationResult:
+    columns = [
+        str(item.get("column"))
+        for item in _mapping_items(context.node.config.get("fields"))
+        if item.get("column")
+    ]
+    return _passthrough_with_columns(
+        context,
+        columns,
+        field_name="fields",
+    )
+
+
+def _quality(context: SchemaPropagationContext) -> SchemaPropagationResult:
+    rules = _mapping_items(context.node.config.get("rules"))
+    diagnostics = _unknown_columns(
+        context.node,
+        context.input_state,
+        [str(rule.get("column") or "") for rule in rules],
+        field_name="rules",
+    )
+    if context.node.config.get("action", "annotate") != "annotate":
+        return SchemaPropagationResult(
+            context.input_state,
+            tuple(diagnostics),
+        )
+    state = _with_column(context.input_state, "_quality_valid", "boolean")
+    state = _with_column(state, "_quality_errors", "array")
+    return SchemaPropagationResult(state, tuple(diagnostics))
+
+
+def _reconcile(context: SchemaPropagationContext) -> SchemaPropagationResult:
+    node = context.node
+    left_state = context.port_state("left")
+    right_state = context.port_state("right")
+    left_compare, right_compare = _comparison_columns(
+        node.config.get("compare_columns")
+    )
+    diagnostics = [
+        *_unknown_columns(
+            node,
+            left_state,
+            node.config.get("left_keys"),
+            field_name="left_keys",
+        ),
+        *_unknown_columns(
+            node,
+            right_state,
+            node.config.get("right_keys"),
+            field_name="right_keys",
+        ),
+        *_unknown_columns(
+            node,
+            left_state,
+            left_compare,
+            field_name="compare_columns",
+        ),
+        *_unknown_columns(
+            node,
+            right_state,
+            right_compare,
+            field_name="compare_columns",
+        ),
+    ]
+    prefix = str(node.config.get("right_prefix") or "observed_")
+    state = SchemaState(
+        left_state.columns
+        | frozenset(
+            f"{prefix}{column}"
+            for column in right_state.columns
+        )
+        | frozenset(
+            (
+                "_reconciliation_status",
+                "_reconciliation_differences",
+            )
+        ),
+        open=left_state.open or right_state.open,
+        types={
+            **left_state.types,
+            **{
+                f"{prefix}{column}": right_state.type_of(column)
+                for column in right_state.columns
+            },
+            "_reconciliation_status": "string",
+            "_reconciliation_differences": "array",
+        },
+    )
+    return SchemaPropagationResult(state, tuple(diagnostics))
+
+
+def _comparison_columns(value: object) -> tuple[list[str], list[str]]:
+    left: list[str] = []
+    right: list[str] = []
+    if not isinstance(value, list):
+        return left, right
+    for item in value:
+        if isinstance(item, str):
+            left.append(item)
+            right.append(item)
+        elif isinstance(item, dict):
+            left_name = str(item.get("left") or item.get("column") or "")
+            right_name = str(
+                item.get("right")
+                or item.get("left")
+                or item.get("column")
+                or ""
+            )
+            left.append(left_name)
+            right.append(right_name)
+    return left, right
+
+
+def _subflow(context: SchemaPropagationContext) -> SchemaPropagationResult:
+    output_schema = _configured_schema(
+        context.node.config.get("output_schema")
+    )
+    return SchemaPropagationResult(
+        output_schema
+        if output_schema.columns
+        else unknown_schema()
+    )
+
+
+def _identity(context: SchemaPropagationContext) -> SchemaPropagationResult:
+    return SchemaPropagationResult(context.input_state)
+
+
+def _passthrough_with_columns(
+    context: SchemaPropagationContext,
+    columns: object,
+    *,
+    field_name: str,
+) -> SchemaPropagationResult:
+    diagnostics = _unknown_columns(
+        context.node,
+        context.input_state,
+        columns,
+        field_name=field_name,
+    )
+    return SchemaPropagationResult(
+        context.input_state,
+        tuple(diagnostics),
+    )
+
+
+def _with_column(
+    state: SchemaState,
+    name: str,
+    data_type: str,
+) -> SchemaState:
+    return SchemaState(
+        state.columns | frozenset((name,)),
+        open=state.open,
+        types={**state.types, name: data_type},
+    )
+
+
+def _configured_schema(value: object) -> SchemaState:
+    if not isinstance(value, list):
+        return unknown_schema()
+    columns = {
+        item
+        if isinstance(item, str)
+        else str(item.get("name"))
+        for item in value
+        if (
+            isinstance(item, str)
+            and item
+            or isinstance(item, dict)
+            and item.get("name")
+        )
+    }
+    types = {
+        str(item["name"]): str(
+            item.get("data_type")
+            or item.get("type")
+            or "unknown"
+        )
+        for item in value
+        if isinstance(item, dict) and item.get("name")
+    }
+    return SchemaState(
+        frozenset(columns),
+        open=not columns,
+        types=types,
+    )
+
+
+def _inline_schema(value: object) -> SchemaState:
+    if not isinstance(value, list):
+        return unknown_schema()
+    rows = [row for row in value if isinstance(row, dict)]
+    columns = {
+        str(column)
+        for row in rows
+        for column in row
+    }
+    types = {
+        column: _observed_column_type(rows, column)
+        for column in columns
+    }
+    return SchemaState(
+        frozenset(columns),
+        open=not columns,
+        types=types,
+    )
+
+
+def _observed_column_type(
+    rows: list[dict[str, Any]],
+    column: str,
+) -> str:
+    observed = {
+        _schema_value_type(row.get(column))
+        for row in rows
+        if row.get(column) is not None
+    }
+    return observed.pop() if len(observed) == 1 else "unknown"
+
+
+def _schema_value_type(value: object) -> str:
+    type_checks = (
+        (bool, "boolean"),
+        (int, "integer"),
+        (float, "number"),
+        (str, "string"),
+        (list, "array"),
+        (dict, "object"),
+    )
+    return next(
+        (
+            name
+            for value_type, name in type_checks
+            if isinstance(value, value_type)
+        ),
+        "unknown",
+    )
+
+
+def _merged_schema_types(
+    states: tuple[SchemaState, ...],
+) -> dict[str, str]:
+    columns = frozenset().union(*(state.columns for state in states))
+    return {
+        column: _merged_column_type(states, column)
+        for column in columns
+    }
+
+
+def _merged_column_type(
+    states: tuple[SchemaState, ...],
+    column: str,
+) -> str:
+    observed = {
+        state.type_of(column)
+        for state in states
+        if (
+            column in state.columns
+            and state.type_of(column) != "unknown"
+        )
+    }
+    return observed.pop() if len(observed) == 1 else "unknown"
+
+
+def _derive_result_type(
+    operation: str,
+    source_types: list[str],
+) -> str:
+    if operation in {"upper", "lower", "trim", "concat"}:
+        return "string"
+    if operation in {"add", "subtract", "multiply", "divide"}:
+        return (
+            "number"
+            if "number" in source_types or operation == "divide"
+            else "integer"
+        )
+    if operation in {"copy", "coalesce"}:
+        concrete = {
+            item
+            for item in source_types
+            if item not in {"unknown", "null"}
+        }
+        return (
+            concrete.pop()
+            if len(concrete) == 1
+            else "unknown"
+        )
+    return "unknown"
+
+
+def _node_expression(
+    node: GraphNode,
+    field_name: str,
+) -> tuple[object | None, list[DataflowDiagnostic]]:
+    try:
+        return (
+            parse_expression(str(node.config.get(field_name) or "")),
+            [],
+        )
+    except ExpressionError as exc:
+        return (
+            None,
+            [
+                _error(
+                    "expression.invalid",
+                    str(exc),
+                    node_id=node.id,
+                    field=field_name,
+                )
+            ],
+        )
+
+
+def _unknown_columns(
+    node: GraphNode,
+    state: SchemaState,
+    columns: object,
+    *,
+    field_name: str,
+) -> list[DataflowDiagnostic]:
+    return [
+        _error(
+            "schema.unknown_column",
+            f"Column {column!r} is not available at this node.",
+            node_id=node.id,
+            field=field_name,
+        )
+        for column in _text_items(columns)
+        if not state.knows(column)
+    ]
+
+
+def _duplicate_outputs(
+    node: GraphNode,
+    columns: list[str],
+    *,
+    field_name: str,
+) -> list[DataflowDiagnostic]:
+    duplicates = sorted(
+        column
+        for column in set(columns)
+        if columns.count(column) > 1
+    )
+    if not duplicates:
+        return []
+    return [
+        _error(
+            "schema.duplicate_output",
+            f"Output column names must be unique: {', '.join(duplicates)}.",
+            node_id=node.id,
+            field=field_name,
+        )
+    ]
+
+
+def _text_items(value: object) -> list[str]:
+    if not isinstance(value, list):
+        return []
+    return [
+        item
+        for item in value
+        if isinstance(item, str) and item
+    ]
+
+
+def _mapping_items(value: object) -> list[dict[str, Any]]:
+    if not isinstance(value, list):
+        return []
+    return [
+        item
+        for item in value
+        if isinstance(item, dict)
+    ]
+
+
+def unknown_schema() -> SchemaState:
+    return SchemaState(frozenset(), open=True)
+
+
+def _error(
+    code: str,
+    message: str,
+    *,
+    node_id: str | None = None,
+    field: str | None = None,
+) -> DataflowDiagnostic:
+    return DataflowDiagnostic(
+        severity="error",
+        code=code,
+        message=message,
+        node_id=node_id,
+        field=field,
+    )
+
+
+def _warning(
+    code: str,
+    message: str,
+    *,
+    node_id: str | None = None,
+    field: str | None = None,
+) -> DataflowDiagnostic:
+    return DataflowDiagnostic(
+        severity="warning",
+        code=code,
+        message=message,
+        node_id=node_id,
+        field=field,
+    )
+
+
+def register_schema_propagators() -> None:
+    propagators = {
+        "source.inline": _inline_source,
+        "source.reference": _reference_source,
+        "combine.union": _union,
+        "combine.join": _join,
+        "filter": _filter,
+        "filter.expression": _filter_expression,
+        "distinct": _distinct,
+        "select": _select,
+        "derive": _derive,
+        "expression": _expression,
+        "convert": _convert_or_replace,
+        "replace": _convert_or_replace,
+        "aggregate": _aggregate,
+        "sort": _sort,
+        "limit": _identity,
+        "quality.rules": _quality,
+        "reconcile.compare": _reconcile,
+        "subflow": _subflow,
+        "output": _identity,
+    }
+    for node_type, propagator in propagators.items():
+        if OPERATOR_REGISTRY.schema_propagator(node_type) is None:
+            OPERATOR_REGISTRY.register_schema_propagator(
+                node_type,
+                propagator,
+            )
+
+
+__all__ = [
+    "SchemaPropagationContext",
+    "SchemaPropagationResult",
+    "SchemaState",
+    "register_schema_propagators",
+    "validate_graph_schemas",
+]

src/govoplan_dataflow/backend/service.py

@@ -33,6 +33,7 @@ from govoplan_dataflow.backend.db.models import (
 from govoplan_dataflow.backend.executor import (
     EXECUTOR_VERSION,
     PipelineExecutionError,
+    PipelineExecutionResult,
     ResolvedSource,
     execute_preview,
 )
@@ -757,6 +758,55 @@ def start_pipeline_run(
     registry: object | None,
     request: DataflowRunRequest,
 ) -> tuple[DataflowRun, bool]:
+    pipeline, revision = _run_definition(
+        session,
+        tenant_id=tenant_id,
+        principal=principal,
+        registry=registry,
+        request=request,
+    )
+    idempotency_key, request_hash = _validated_run_identity(request)
+    existing = _existing_pipeline_run(
+        session,
+        tenant_id=tenant_id,
+        pipeline_id=pipeline.id,
+        idempotency_key=idempotency_key,
+        request_hash=request_hash,
+    )
+    if existing is not None:
+        return existing, True
+    run = _new_pipeline_run(
+        tenant_id=tenant_id,
+        actor_id=actor_id,
+        pipeline=pipeline,
+        revision=revision,
+        request=request,
+        idempotency_key=idempotency_key,
+        request_hash=request_hash,
+    )
+    session.add(run)
+    session.flush()
+    _execute_pipeline_run(
+        session,
+        run=run,
+        pipeline=pipeline,
+        revision=revision,
+        request=request,
+        principal=principal,
+        registry=registry,
+    )
+    session.flush()
+    return run, False
+
+
+def _run_definition(
+    session: Session,
+    *,
+    tenant_id: str,
+    principal: ApiPrincipal,
+    registry: object | None,
+    request: DataflowRunRequest,
+) -> tuple[DataflowPipeline, DataflowPipelineRevision]:
     pipeline_id = _strip_ref(request.pipeline_ref, "pipeline:")
     if not pipeline_id:
         raise DataflowNotFoundError("Dataflow pipeline not found")
@@ -784,6 +834,10 @@ def start_pipeline_run(
         pipeline=pipeline,
         revision=request.revision,
     )
+    return pipeline, revision
+
+
+def _validated_run_identity(request: DataflowRunRequest) -> tuple[str, str]:
     idempotency_key = request.idempotency_key.strip()
     if not idempotency_key or len(idempotency_key) > 255:
         raise DataflowConflictError(
@@ -793,11 +847,21 @@ def start_pipeline_run(
         raise DataflowConflictError(
             "The bounded Dataflow runner supports between 1 and 500 output rows."
         )
-    request_hash = _run_request_hash(request)
+    return idempotency_key, _run_request_hash(request)
+
+
+def _existing_pipeline_run(
+    session: Session,
+    *,
+    tenant_id: str,
+    pipeline_id: str,
+    idempotency_key: str,
+    request_hash: str,
+) -> DataflowRun | None:
     existing = session.scalar(
         select(DataflowRun).where(
             DataflowRun.tenant_id == tenant_id,
-            DataflowRun.pipeline_id == pipeline.id,
+            DataflowRun.pipeline_id == pipeline_id,
             DataflowRun.idempotency_key == idempotency_key,
         )
     )
@@ -807,11 +871,20 @@ def start_pipeline_run(
                 "The Dataflow run idempotency key was already used with "
                 "different parameters."
             )
-        return existing, True
+    return existing
 
-    graph = PipelineGraph.model_validate(revision.graph)
-    started_at = utcnow()
-    run = DataflowRun(
+
+def _new_pipeline_run(
+    *,
+    tenant_id: str,
+    actor_id: str | None,
+    pipeline: DataflowPipeline,
+    revision: DataflowPipelineRevision,
+    request: DataflowRunRequest,
+    idempotency_key: str,
+    request_hash: str,
+) -> DataflowRun:
+    return DataflowRun(
         tenant_id=tenant_id,
         pipeline_id=pipeline.id,
         pipeline_revision_id=revision.id,
@@ -838,15 +911,24 @@ def start_pipeline_run(
         diagnostics=[],
         input_row_count=0,
         output_row_count=0,
-        started_at=started_at,
+        started_at=utcnow(),
         created_by=actor_id,
     )
-    session.add(run)
-    session.flush()
 
+
+def _execute_pipeline_run(
+    session: Session,
+    *,
+    run: DataflowRun,
+    pipeline: DataflowPipeline,
+    revision: DataflowPipelineRevision,
+    request: DataflowRunRequest,
+    principal: ApiPrincipal,
+    registry: object | None,
+) -> None:
     try:
         result = execute_preview(
-            graph,
+            PipelineGraph.model_validate(revision.graph),
             row_limit=request.row_limit,
             source_resolver=_datasource_source_resolver(
                 session=session,
@@ -854,91 +936,128 @@ def start_pipeline_run(
                 registry=registry,
             ),
         )
-        if request.publication and (
-            result.truncated
-            or any(
-                bool(item.get("truncated"))
-                for item in result.source_fingerprints
-            )
-        ):
-            raise PipelineExecutionError(
-                "The bounded runner cannot publish a truncated result or a "
-                "result calculated from truncated source data."
-            )
-
-        run.source_fingerprints = result.source_fingerprints
-        run.result_schema = [
-            item.model_dump(mode="json") for item in result.columns
-        ]
-        run.diagnostics = [
-            item.model_dump(mode="json") for item in result.diagnostics
-        ]
-        run.input_row_count = result.input_row_count
-        run.output_row_count = result.total_rows
-
+        _apply_pipeline_result(run, result)
         if request.publication:
-            publisher = datasource_publication(registry)
-            if publisher is None:
-                raise PipelineExecutionError(
-                    "Publishing Dataflow output requires the Datasources "
-                    "publication capability."
-                )
-            target = request.publication
-            publication = publisher.publish_rows(
+            _ensure_publishable(result)
+            _publish_pipeline_result(
                 session,
-                principal,
-                request=DatasourcePublicationRequest(
-                    producer_module="dataflow",
-                    producer_run_ref=f"dataflow-run:{run.id}",
-                    idempotency_key=f"{pipeline.id}:{idempotency_key}",
-                    rows=tuple(dict(row) for row in result.rows),
-                    target_datasource_ref=target.target_datasource_ref,
-                    name=target.name or f"{pipeline.name} output",
-                    source_name=target.source_name,
-                    description=target.description,
-                    freeze=target.freeze,
-                    frozen_label=target.frozen_label,
-                    set_current=target.set_current,
-                    provenance={
-                        "pipeline_ref": f"pipeline:{pipeline.id}",
-                        "pipeline_revision": revision.revision,
-                        "definition_hash": revision.content_hash,
-                        "source_fingerprints": result.source_fingerprints,
-                    },
-                    metadata={
-                        **dict(target.metadata),
-                        "dataflow_run_ref": f"dataflow-run:{run.id}",
-                    },
-                ),
+                run=run,
+                pipeline=pipeline,
+                revision=revision,
+                request=request,
+                result=result,
+                principal=principal,
+                registry=registry,
             )
-            run.output_publication_ref = publication.ref
-            run.output_datasource_ref = publication.datasource.ref
-            run.output_materialization_ref = publication.materialization.ref
         run.status = "succeeded"
         run.finished_at = utcnow()
         run.error = None
     except (DatasourceError, PipelineExecutionError) as exc:
-        run.status = "failed"
-        run.finished_at = utcnow()
-        run.error = str(exc)
-        diagnostics = list(getattr(exc, "diagnostics", ()))
-        diagnostics.append(
-            DataflowDiagnostic(
-                severity="error",
-                code="run.execution",
-                message=str(exc),
-                node_id=getattr(exc, "node_id", None),
-            )
+        _mark_pipeline_run_failed(run, exc)
+
+
+def _apply_pipeline_result(
+    run: DataflowRun,
+    result: PipelineExecutionResult,
+) -> None:
+    run.source_fingerprints = result.source_fingerprints
+    run.result_schema = [
+        item.model_dump(mode="json") for item in result.columns
+    ]
+    run.diagnostics = [
+        item.model_dump(mode="json") for item in result.diagnostics
+    ]
+    run.input_row_count = result.input_row_count
+    run.output_row_count = result.total_rows
+
+
+def _ensure_publishable(result: PipelineExecutionResult) -> None:
+    source_truncated = any(
+        bool(item.get("truncated"))
+        for item in result.source_fingerprints
+    )
+    if result.truncated or source_truncated:
+        raise PipelineExecutionError(
+            "The bounded runner cannot publish a truncated result or a "
+            "result calculated from truncated source data."
         )
-        run.diagnostics = [
-            item.model_dump(mode="json") for item in diagnostics
-        ]
-        run.source_fingerprints = list(
-            getattr(exc, "source_fingerprints", ())
+
+
+def _publish_pipeline_result(
+    session: Session,
+    *,
+    run: DataflowRun,
+    pipeline: DataflowPipeline,
+    revision: DataflowPipelineRevision,
+    request: DataflowRunRequest,
+    result: PipelineExecutionResult,
+    principal: ApiPrincipal,
+    registry: object | None,
+) -> None:
+    publisher = datasource_publication(registry)
+    if publisher is None:
+        raise PipelineExecutionError(
+            "Publishing Dataflow output requires the Datasources "
+            "publication capability."
         )
-        run.input_row_count = int(getattr(exc, "input_row_count", 0))
-    session.flush()
-    return run, False
+    target = request.publication
+    if target is None:
+        return
+    publication = publisher.publish_rows(
+        session,
+        principal,
+        request=DatasourcePublicationRequest(
+            producer_module="dataflow",
+            producer_run_ref=f"dataflow-run:{run.id}",
+            idempotency_key=f"{pipeline.id}:{request.idempotency_key.strip()}",
+            rows=tuple(dict(row) for row in result.rows),
+            target_datasource_ref=target.target_datasource_ref,
+            name=target.name or f"{pipeline.name} output",
+            source_name=target.source_name,
+            description=target.description,
+            freeze=target.freeze,
+            frozen_label=target.frozen_label,
+            set_current=target.set_current,
+            provenance={
+                "pipeline_ref": f"pipeline:{pipeline.id}",
+                "pipeline_revision": revision.revision,
+                "definition_hash": revision.content_hash,
+                "source_fingerprints": result.source_fingerprints,
+            },
+            metadata={
+                **dict(target.metadata),
+                "dataflow_run_ref": f"dataflow-run:{run.id}",
+            },
+        ),
+    )
+    run.output_publication_ref = publication.ref
+    run.output_datasource_ref = publication.datasource.ref
+    run.output_materialization_ref = publication.materialization.ref
+
+
+def _mark_pipeline_run_failed(
+    run: DataflowRun,
+    exc: DatasourceError | PipelineExecutionError,
+) -> None:
+    run.status = "failed"
+    run.finished_at = utcnow()
+    run.error = str(exc)
+    diagnostics = list(getattr(exc, "diagnostics", ()))
+    diagnostics.append(
+        DataflowDiagnostic(
+            severity="error",
+            code="run.execution",
+            message=str(exc),
+            node_id=getattr(exc, "node_id", None),
+        )
+    )
+    run.diagnostics = [
+        item.model_dump(mode="json") for item in diagnostics
+    ]
+    run.source_fingerprints = list(
+        getattr(exc, "source_fingerprints", ())
+    )
+    run.input_row_count = int(getattr(exc, "input_row_count", 0))
 
 
 def cancel_pipeline_run(