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)