991 lines
27 KiB
Python
991 lines
27 KiB
Python
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]:
|
|
_, diagnostics = propagate_graph_schemas(graph, ordered=ordered)
|
|
return diagnostics
|
|
|
|
|
|
def propagate_graph_schemas(
|
|
graph: PipelineGraph,
|
|
*,
|
|
ordered: list[str],
|
|
) -> tuple[dict[str, SchemaState], 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 schemas, 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",
|
|
),
|
|
]
|
|
if node.config.get("join_type", "inner") in {"semi", "anti"}:
|
|
return SchemaPropagationResult(left_state, tuple(diagnostics))
|
|
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 _calculate(context: SchemaPropagationContext) -> SchemaPropagationResult:
|
|
node = context.node
|
|
state = context.input_state
|
|
diagnostics: list[DataflowDiagnostic] = []
|
|
for item in _mapping_items(node.config.get("calculations")):
|
|
source = str(item.get("expression") or "")
|
|
target = str(item.get("target_column") or "")
|
|
try:
|
|
parsed = parse_expression(source)
|
|
except ExpressionError as exc:
|
|
diagnostics.append(
|
|
_error(
|
|
"expression.invalid",
|
|
str(exc),
|
|
node_id=node.id,
|
|
field="calculations",
|
|
)
|
|
)
|
|
continue
|
|
diagnostics.extend(
|
|
_unknown_columns(
|
|
node,
|
|
state,
|
|
list(parsed.columns),
|
|
field_name="calculations",
|
|
)
|
|
)
|
|
inferred = infer_expression_type(
|
|
parsed,
|
|
{
|
|
name: state.type_of(name) # type: ignore[dict-item]
|
|
for name in state.columns
|
|
},
|
|
)
|
|
expected = str(item.get("result_type") or "unknown")
|
|
if expected != "unknown" and inferred not in {
|
|
"unknown",
|
|
"null",
|
|
expected,
|
|
}:
|
|
diagnostics.append(
|
|
_warning(
|
|
"calculate.type_mismatch",
|
|
(
|
|
f"Calculation for {target!r} infers {inferred}, "
|
|
f"not {expected}."
|
|
),
|
|
node_id=node.id,
|
|
field="calculations",
|
|
)
|
|
)
|
|
if target:
|
|
state = _with_column(
|
|
state,
|
|
target,
|
|
expected if expected != "unknown" else inferred,
|
|
)
|
|
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 _rank(context: SchemaPropagationContext) -> SchemaPropagationResult:
|
|
node = context.node
|
|
columns = [
|
|
*_text_items(node.config.get("partition_by")),
|
|
*[
|
|
str(item.get("column"))
|
|
for item in _mapping_items(node.config.get("order_by"))
|
|
if item.get("column")
|
|
],
|
|
]
|
|
diagnostics = _unknown_columns(
|
|
node,
|
|
context.input_state,
|
|
columns,
|
|
field_name="order_by",
|
|
)
|
|
target = str(node.config.get("target_column") or "")
|
|
state = (
|
|
_with_column(context.input_state, target, "integer")
|
|
if target
|
|
else context.input_state
|
|
)
|
|
return SchemaPropagationResult(state, tuple(diagnostics))
|
|
|
|
|
|
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,
|
|
"calculate": _calculate,
|
|
"convert": _convert_or_replace,
|
|
"replace": _convert_or_replace,
|
|
"aggregate": _aggregate,
|
|
"sort": _sort,
|
|
"window.rank": _rank,
|
|
"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",
|
|
]
|