Refactor dataflow operators around runtime registry

This commit is contained in:
2026-07-29 16:27:26 +02:00
parent 946202ef01
commit 69509d5cc2
7 changed files with 3660 additions and 2075 deletions
+381 -113
View File
@@ -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,84 +112,155 @@ 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]
def _execute_preview_node(
plan: _PreviewPlan,
state: _PreviewState,
node: GraphNode,
*,
source_resolver: SourceResolver | None,
execution_depth: int,
) -> None:
node_started = time.monotonic()
node_inputs = inputs.get(node_id, {})
node_inputs = plan.inputs_by_port.get(node.id, {})
input_sets = [
outputs[source_id]
state.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),
)
_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,
)
)
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,
)
execution = executor(
return executor(
OperatorExecutionContext(
node=node,
inputs_by_port=node_inputs,
outputs=outputs,
input_sets=input_sets,
input_rows=input_rows,
input_rows=input_sets[0] if len(input_sets) == 1 else [],
source_resolver=source_resolver,
execution_depth=_execution_depth,
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(
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",
@@ -178,12 +269,63 @@ def execute_preview(
)
for message in execution.messages
)
if len(json.dumps(output_rows, default=str).encode("utf-8")) > MAX_RESULT_BYTES:
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,
node_id=node_id,
)
except (PipelineExecutionError, ArithmeticError, KeyError, TypeError, ValueError) as exc:
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)
@@ -192,51 +334,93 @@ def execute_preview(
failed = NodePreviewDiagnostic(
node_id=execution_error.node_id or node.id,
status="failed",
input_rows=sum(len(rows) for rows in input_sets),
input_rows=_input_row_total(input_sets),
output_rows=0,
duration_ms=round((time.monotonic() - node_started) * 1000, 3),
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([*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),
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
outputs[node.id] = output_rows
node_diagnostics.append(
NodePreviewDiagnostic(
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=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,
)
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),
)
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 _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":
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)
if operation == "coalesce":
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)}
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:
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]
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":
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
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
output.append(result)
return output
_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)
+347 -130
View File
@@ -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,81 +142,114 @@ def infer_expression_type(
def _evaluate(expression: exp.Expression, row: dict[str, Any]) -> Any:
if isinstance(expression, exp.Paren):
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)
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)
def _evaluate_negation(expression: exp.Expression, row: dict[str, Any]) -> Any:
value = _evaluate_child(expression, row)
return None if value is None else -value
if isinstance(expression, exp.Not):
return not bool(_evaluate(expression.this, row))
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):
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
if isinstance(expression, (exp.Add, exp.Sub, exp.Mul, exp.Div, exp.Mod)):
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 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)):
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)
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 _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)
if isinstance(expression, exp.Abs):
value = _evaluate(expression.this, row)
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)
if isinstance(expression, exp.Round):
value = _evaluate(expression.this, row)
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)
if isinstance(expression, exp.Coalesce):
def _evaluate_coalesce(expression: exp.Expression, row: dict[str, Any]) -> Any:
arguments = [expression.this, *expression.expressions]
return next(
(
@@ -225,47 +259,136 @@ def _evaluate(expression: exp.Expression, row: dict[str, Any]) -> Any:
),
None,
)
if isinstance(expression, exp.Concat):
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
)
if isinstance(expression, exp.Replace):
value = _evaluate(expression.this, row)
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))
if isinstance(expression, exp.Substring):
value = _evaluate(expression.this, row)
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")
return (
str(value)[start:]
if length is None
else str(value)[start : start + int(_evaluate(length, row))]
)
if isinstance(expression, exp.Cast):
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(expression.this, row),
_cast_type(expression),
_evaluate_child(expression, row),
_cast_type(expression), # type: ignore[arg-type]
on_error="fail",
)
if isinstance(expression, exp.Case):
base = _evaluate(expression.this, row) if expression.this is not None else None
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 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:
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
raise ExpressionError(
f"Expression element {type(expression).__name__} cannot be evaluated."
)
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(
@@ -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:
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(expression.this) else "integer"
if isinstance(
expression,
(
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:
File diff suppressed because it is too large Load Diff
@@ -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",
]
+147 -28
View File
@@ -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,18 +936,30 @@ 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."
_apply_pipeline_result(run, result)
if request.publication:
_ensure_publishable(result)
_publish_pipeline_result(
session,
run=run,
pipeline=pipeline,
revision=revision,
request=request,
result=result,
principal=principal,
registry=registry,
)
run.status = "succeeded"
run.finished_at = utcnow()
run.error = None
except (DatasourceError, PipelineExecutionError) as exc:
_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
@@ -876,7 +970,30 @@ def start_pipeline_run(
run.input_row_count = result.input_row_count
run.output_row_count = result.total_rows
if request.publication:
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."
)
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(
@@ -884,13 +1001,15 @@ def start_pipeline_run(
"publication capability."
)
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}:{idempotency_key}",
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",
@@ -914,10 +1033,12 @@ def start_pipeline_run(
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:
def _mark_pipeline_run_failed(
run: DataflowRun,
exc: DatasourceError | PipelineExecutionError,
) -> None:
run.status = "failed"
run.finished_at = utcnow()
run.error = str(exc)
@@ -937,8 +1058,6 @@ def start_pipeline_run(
getattr(exc, "source_fingerprints", ())
)
run.input_row_count = int(getattr(exc, "input_row_count", 0))
session.flush()
return run, False
def cancel_pipeline_run(
File diff suppressed because it is too large Load Diff
+49
View File
@@ -276,9 +276,58 @@ class DataflowOperatorTests(unittest.TestCase):
self.assertEqual(set(NODE_TYPES), set(coverage))
for node_type, facets in coverage.items():
self.assertIn("config_validator", facets, node_type)
self.assertIn("schema_propagator", facets, node_type)
self.assertIn("executor", facets, node_type)
self.assertIn("sql_renderer", facets, node_type)
def test_maximum_size_linear_graph_validates_and_executes(self) -> None:
nodes = [
node(
"source",
"source.inline",
{
"source_name": "bounded_input",
"rows": [{"value": 1}],
},
x=0,
)
]
edges: list[GraphEdge] = []
previous_id = "source"
for index in range(1, 99):
node_id = f"limit-{index}"
nodes.append(
node(
node_id,
"limit",
{"count": 1},
x=index * 10,
)
)
edges.append(
GraphEdge(
id=f"edge-{index}",
source=previous_id,
target=node_id,
)
)
previous_id = node_id
nodes.append(node("output", "output", {}, x=990))
edges.append(
GraphEdge(
id="edge-99",
source=previous_id,
target="output",
)
)
graph = PipelineGraph(nodes=nodes, edges=edges)
self.assertEqual([], validate_graph(graph))
result = execute_preview(graph, row_limit=1)
self.assertEqual([{"value": 1}], result.rows)
self.assertEqual(100, len(result.node_diagnostics))
def test_expression_node_renders_through_registered_sql_compiler(self) -> None:
graph = PipelineGraph(
nodes=[