Add typed Arrow execution backend boundary

This commit is contained in:
2026-07-29 16:47:14 +02:00
parent 69509d5cc2
commit 8a11d7571b
17 changed files with 2264 additions and 49 deletions
+20 -4
View File
@@ -59,10 +59,26 @@ expression filters and calculated columns. It supports literals, columns,
arithmetic, comparisons, boolean logic, `CASE`, safe casts, and a bounded arithmetic, comparisons, boolean logic, `CASE`, safe casts, and a bounded
string/numeric function catalogue without Python evaluation or effectful SQL. string/numeric function catalogue without Python evaluation or effectful SQL.
Node definitions, validators, preview executors, and SQL renderers are Node definitions, validators, schema propagators, preview executors, and SQL
registered independently in the operator registry. Adding a node no longer renderers are registered independently in the operator registry. Adding a node
requires another branch in the preview or graph-to-SQL dispatch loop. Operators no longer requires another branch in the preview or graph-to-SQL dispatch loop.
that cannot be represented by constrained SQL declare that explicitly. Operators that cannot be represented by constrained SQL declare that
explicitly.
Every executable graph also has a versioned typed IR. It preserves graph
identity and layout while giving ports, schemas, expressions, parameters,
lineage, diagnostics, semantic hashes, and physical results stable contracts.
The backend boundary exchanges bounded typed columnar batches that can be
serialized as Arrow IPC. The deterministic Python executor remains the default
and fallback.
Install `.[analytics]` to enable the analytical backend. It executes only SQL
generated from validated graphs in a separate short-lived DuckDB process.
Inputs and outputs cross the process boundary as Arrow IPC; external access,
extension installation/loading, persistent secrets, configuration changes,
temporary spill files, and multiple DuckDB threads are disabled. Wall-clock,
row, byte, memory, file-descriptor, and backend-concurrency limits are applied.
Callers can explicitly request `reference`, `duckdb`, or `auto` for a preview.
Preview reads at most 250 rows per source and enforces time, intermediate-row, Preview reads at most 250 rows per source and enforces time, intermediate-row,
result-byte, graph-node, and response-row bounds. Saved previews record the result-byte, graph-node, and response-row bounds. Saved previews record the
+6
View File
@@ -15,6 +15,12 @@ dependencies = [
"sqlglot>=30.14,<31", "sqlglot>=30.14,<31",
] ]
[project.optional-dependencies]
analytics = [
"duckdb>=1.5.5,<2",
"pyarrow>=24,<26",
]
[tool.setuptools.packages.find] [tool.setuptools.packages.find]
where = ["src"] where = ["src"]
@@ -0,0 +1,25 @@
from govoplan_dataflow.backend.backends.base import (
BackendExecutionError,
BackendExecutionRequest,
BackendExecutionResult,
BackendSource,
ExecutionBackend,
ExecutionBudget,
)
from govoplan_dataflow.backend.backends.registry import (
EXECUTION_BACKENDS,
ExecutionBackendRegistry,
execute_typed_graph,
)
__all__ = [
"BackendExecutionError",
"BackendExecutionRequest",
"BackendExecutionResult",
"BackendSource",
"EXECUTION_BACKENDS",
"ExecutionBackend",
"ExecutionBackendRegistry",
"ExecutionBudget",
"execute_typed_graph",
]
@@ -0,0 +1,196 @@
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any, Mapping, Protocol, runtime_checkable
from govoplan_dataflow.backend.batches import TypedBatch
from govoplan_dataflow.backend.ir import (
IrExecutionResult,
IrField,
IrLineage,
IrSchema,
TypedGraphIr,
schema_from_fields,
)
from govoplan_dataflow.backend.planner import ExecutionPlan
from govoplan_dataflow.backend.schemas import (
DataflowDiagnostic,
NodePreviewDiagnostic,
NodePreviewResult,
)
@dataclass(frozen=True, slots=True)
class ExecutionBudget:
max_output_rows: int = 500
max_batch_bytes: int = 1_000_000
max_wall_seconds: float = 2.0
max_memory_bytes: int = 256 * 1024 * 1024
max_concurrency: int = 1
def __post_init__(self) -> None:
if self.max_output_rows < 1:
raise ValueError("Execution output row limit must be positive.")
if self.max_batch_bytes < 1:
raise ValueError("Execution byte limit must be positive.")
if self.max_wall_seconds <= 0:
raise ValueError("Execution time limit must be positive.")
if self.max_memory_bytes < 64 * 1024 * 1024:
raise ValueError("Execution memory limit must be at least 64 MiB.")
if self.max_concurrency < 1:
raise ValueError("Execution concurrency limit must be positive.")
@dataclass(frozen=True, slots=True)
class BackendSource:
node_id: str
batch: TypedBatch
source_ref: str
provider: str
fingerprint: str
total_rows: int
truncated: bool = False
source_name: str | None = None
kind: str = "datasource"
def lineage_payload(self) -> dict[str, Any]:
payload: dict[str, Any] = {
"node_id": self.node_id,
"source_name": self.source_name,
"kind": self.kind,
"fingerprint": self.fingerprint,
"row_count": self.total_rows,
}
if self.kind != "inline":
payload.update(
{
"source_ref": self.source_ref,
"provider": self.provider,
"preview_rows": self.batch.row_count,
"truncated": self.truncated,
}
)
return payload
@dataclass(frozen=True, slots=True)
class BackendExecutionRequest:
plan: ExecutionPlan
sources: Mapping[str, BackendSource] = field(default_factory=dict)
budget: ExecutionBudget = field(default_factory=ExecutionBudget)
preview_node_id: str | None = None
@property
def graph(self) -> TypedGraphIr:
return self.plan.graph
@dataclass(frozen=True, slots=True)
class BackendExecutionResult:
contract: IrExecutionResult
batch: TypedBatch
node_diagnostics: tuple[NodePreviewDiagnostic, ...] = ()
node_preview: NodePreviewResult | None = None
metadata: Mapping[str, Any] = field(default_factory=dict)
@property
def rows(self) -> list[dict[str, Any]]:
return self.batch.to_rows()
class BackendExecutionError(RuntimeError):
def __init__(
self,
message: str,
*,
code: str = "backend.execution",
node_id: str | None = None,
diagnostics: tuple[DataflowDiagnostic, ...] = (),
) -> None:
super().__init__(message)
self.code = code
self.node_id = node_id
self.diagnostics = diagnostics
@runtime_checkable
class ExecutionBackend(Protocol):
name: str
version: str
def available(self) -> bool:
...
def supports(self, request: BackendExecutionRequest) -> bool:
...
def execute(
self,
request: BackendExecutionRequest,
) -> BackendExecutionResult:
...
def request_lineage(request: BackendExecutionRequest) -> IrLineage:
fingerprints = tuple(
source.lineage_payload()
for source in request.sources.values()
)
return request.graph.lineage.model_copy(
update={"source_fingerprints": fingerprints}
)
def canonical_result_schema(
request: BackendExecutionRequest,
observed: IrSchema,
) -> IrSchema:
output = next(
(
node.output_schema
for node in request.graph.nodes
if node.type == "output"
),
None,
)
if output is None:
return observed
declared = {
field.name: field
for field in output.fields
}
fields = tuple(
_resolved_result_field(field, declared.get(field.name))
for field in observed.fields
)
return schema_from_fields(fields, open=output.open)
def _resolved_result_field(
observed: IrField,
declared: IrField | None,
) -> IrField:
if declared is None:
return observed
return declared.model_copy(
update={
"type": (
observed.type
if declared.type in {"unknown", "null"}
else declared.type
),
"nullable": declared.nullable,
}
)
__all__ = [
"BackendExecutionError",
"BackendExecutionRequest",
"BackendExecutionResult",
"BackendSource",
"ExecutionBackend",
"ExecutionBudget",
"canonical_result_schema",
"request_lineage",
]
@@ -0,0 +1,412 @@
from __future__ import annotations
import hashlib
import importlib.util
import json
import multiprocessing
import os
import signal
from typing import Any
import sqlglot
from sqlglot import exp
from govoplan_dataflow.backend.backends.base import (
BackendExecutionError,
BackendExecutionRequest,
BackendExecutionResult,
BackendSource,
canonical_result_schema,
)
from govoplan_dataflow.backend.batches import (
ArrowDependencyError,
TypedBatch,
)
from govoplan_dataflow.backend.ir import (
IrExecutionResult,
IrLineage,
ir_to_graph,
)
from govoplan_dataflow.backend.schemas import GraphNode
DUCKDB_BACKEND_VERSION = "duckdb-isolated-v1"
class DuckDbExecutionBackend:
name = "duckdb"
version = DUCKDB_BACKEND_VERSION
def available(self) -> bool:
return _analytics_available()
def supports(self, request: BackendExecutionRequest) -> bool:
return (
self.available()
and _previews_final_output(request)
and request.plan.generated_sql is not None
)
def execute(
self,
request: BackendExecutionRequest,
) -> BackendExecutionResult:
if not self.available():
raise BackendExecutionError(
"DuckDB execution requires the govoplan-dataflow analytics extra.",
code="backend.unavailable",
)
if not _previews_final_output(request):
raise BackendExecutionError(
"DuckDB execution currently produces the final output only.",
code="backend.preview_node",
node_id=request.preview_node_id,
)
sql_text = request.plan.generated_sql
if sql_text is None:
raise BackendExecutionError(
"This graph cannot be executed by the DuckDB backend.",
code="backend.unsupported_graph",
diagnostics=request.plan.sql_diagnostics,
)
graph = ir_to_graph(request.graph)
_validate_generated_select(sql_text)
sources = _source_batches(request, graph.nodes)
worker_payload = {
"sql": sql_text,
"sources": {
name: source.batch.to_arrow_ipc()
for name, source in sources.items()
},
"max_rows": request.budget.max_output_rows,
"max_bytes": request.budget.max_batch_bytes,
"max_memory_bytes": request.budget.max_memory_bytes,
}
worker_result = _run_worker(
worker_payload,
timeout=request.budget.max_wall_seconds,
)
try:
observed_batch = TypedBatch.from_arrow_ipc(worker_result["ipc"])
batch = TypedBatch.from_rows(
observed_batch.to_rows(),
schema=canonical_result_schema(
request,
observed_batch.schema,
),
)
except (ArrowDependencyError, KeyError, TypeError, ValueError) as exc:
raise BackendExecutionError(
"DuckDB returned an invalid Arrow result.",
code="backend.result_invalid",
) from exc
batch.ensure_within(
max_rows=request.budget.max_output_rows,
max_bytes=request.budget.max_batch_bytes,
)
lineage = _source_lineage(request, sources)
contract = IrExecutionResult(
backend=self.name,
backend_version=f"{self.version}+{worker_result['duckdb_version']}",
graph_semantic_hash=request.graph.semantic_hash,
result_schema=batch.schema,
row_count=int(worker_result["total_rows"]),
byte_count=batch.byte_count,
truncated=bool(worker_result["truncated"]),
diagnostics=request.plan.diagnostics,
lineage=lineage,
)
return BackendExecutionResult(
contract=contract,
batch=batch,
metadata={
"source_fingerprints": tuple(
source.lineage_payload()
for source in sources.values()
),
"isolated_process": True,
"external_access": False,
},
)
def _analytics_available() -> bool:
return (
importlib.util.find_spec("duckdb") is not None
and importlib.util.find_spec("pyarrow") is not None
)
def _previews_final_output(request: BackendExecutionRequest) -> bool:
if request.preview_node_id is None:
return True
return any(
node.id == request.preview_node_id and node.type == "output"
for node in request.graph.nodes
)
def _validate_generated_select(sql_text: str) -> None:
statements = sqlglot.parse(sql_text, read="duckdb")
if len(statements) != 1 or not isinstance(
statements[0],
(exp.Select, exp.Union),
):
raise BackendExecutionError(
"DuckDB accepts one generated SELECT statement only.",
code="backend.sql_rejected",
)
prohibited = (
exp.Command,
exp.Copy,
exp.Create,
exp.Delete,
exp.Drop,
exp.Insert,
exp.Merge,
exp.Update,
)
if any(isinstance(item, prohibited) for item in statements[0].walk()):
raise BackendExecutionError(
"DuckDB rejected an effectful statement.",
code="backend.sql_rejected",
)
def _source_batches(
request: BackendExecutionRequest,
nodes: list[GraphNode],
) -> dict[str, BackendSource]:
sources: dict[str, BackendSource] = {}
for node in nodes:
if not node.type.startswith("source."):
continue
source_name = str(node.config.get("source_name") or "").strip()
source = (
_inline_source(node)
if node.type == "source.inline"
else _request_source(request, node)
)
try:
source.batch.ensure_within(
max_rows=10_000,
max_bytes=request.budget.max_batch_bytes,
)
except ValueError as exc:
raise BackendExecutionError(
str(exc),
code="backend.source_budget",
node_id=node.id,
) from exc
sources[source_name] = source
return sources
def _inline_source(node: GraphNode) -> BackendSource:
rows = [dict(row) for row in node.config.get("rows") or []]
batch = TypedBatch.from_rows(rows)
fingerprint = hashlib.sha256(
json.dumps(
rows,
sort_keys=True,
separators=(",", ":"),
default=str,
).encode("utf-8")
).hexdigest()
return BackendSource(
node_id=node.id,
batch=batch,
source_ref=f"inline:{node.id}",
provider="inline",
fingerprint=fingerprint,
total_rows=batch.row_count,
source_name=str(node.config.get("source_name") or ""),
kind="inline",
)
def _request_source(
request: BackendExecutionRequest,
node: GraphNode,
) -> BackendSource:
keys = (
node.id,
str(node.config.get("source_ref") or ""),
str(node.config.get("source_name") or ""),
)
source = next(
(
request.sources[key]
for key in keys
if key and key in request.sources
),
None,
)
if source is None:
raise BackendExecutionError(
f"No typed source batch was supplied for node {node.id!r}.",
code="backend.source_missing",
node_id=node.id,
)
return source
def _source_lineage(
request: BackendExecutionRequest,
sources: dict[str, BackendSource],
) -> IrLineage:
return request.graph.lineage.model_copy(
update={
"source_fingerprints": tuple(
source.lineage_payload()
for source in sources.values()
)
}
)
def _run_worker(payload: dict[str, Any], *, timeout: float) -> dict[str, Any]:
context = multiprocessing.get_context("spawn")
receive, send = context.Pipe(duplex=False)
process = context.Process(
target=_duckdb_worker,
args=(send, payload),
name="govoplan-dataflow-duckdb",
daemon=True,
)
process.start()
send.close()
try:
if not receive.poll(timeout):
_terminate_worker(process)
raise BackendExecutionError(
"DuckDB execution exceeded its wall-clock budget.",
code="backend.timeout",
)
result = receive.recv()
except EOFError as exc:
process.join(timeout=0.2)
raise BackendExecutionError(
"DuckDB worker exited without a result.",
code="backend.worker_exit",
) from exc
finally:
receive.close()
process.join(timeout=0.2)
if process.is_alive():
_terminate_worker(process)
if not isinstance(result, dict) or not result.get("ok"):
message = (
str(result.get("error"))
if isinstance(result, dict)
else "DuckDB worker returned an invalid response."
)
raise BackendExecutionError(
message,
code="backend.duckdb",
)
return result
def _terminate_worker(process: multiprocessing.Process) -> None:
process.terminate()
process.join(timeout=0.5)
if process.is_alive() and process.pid:
os.kill(process.pid, signal.SIGKILL)
process.join(timeout=0.5)
def _duckdb_worker(send: Any, payload: dict[str, Any]) -> None:
try:
_set_worker_limits(float(payload["max_memory_bytes"]))
result = _execute_duckdb_payload(payload)
send.send({"ok": True, **result})
except BaseException as exc:
send.send(
{
"ok": False,
"error": f"{type(exc).__name__}: {exc}",
}
)
finally:
send.close()
def _set_worker_limits(max_memory_bytes: float) -> None:
try:
import resource
except ImportError:
return
resource.setrlimit(resource.RLIMIT_CORE, (0, 0))
resource.setrlimit(resource.RLIMIT_FSIZE, (0, 0))
resource.setrlimit(resource.RLIMIT_NOFILE, (64, 64))
cpu_seconds = 5
resource.setrlimit(resource.RLIMIT_CPU, (cpu_seconds, cpu_seconds))
os.environ["GOVOPLAN_DATAFLOW_MEMORY_BUDGET"] = str(int(max_memory_bytes))
def _execute_duckdb_payload(payload: dict[str, Any]) -> dict[str, Any]:
import secrets
import duckdb
import pyarrow as pa
memory_bytes = int(payload["max_memory_bytes"])
connection = duckdb.connect(
database=":memory:",
config={
"allow_community_extensions": "false",
"allow_unsigned_extensions": "false",
"autoinstall_known_extensions": "false",
"autoload_known_extensions": "false",
"allow_persistent_secrets": "false",
"memory_limit": f"{memory_bytes}B",
"max_temp_directory_size": "0B",
"threads": "1",
},
)
try:
for name, ipc_payload in payload["sources"].items():
with pa.ipc.open_stream(ipc_payload) as reader:
connection.register(name, reader.read_all())
connection.execute("SET enable_external_access = false")
connection.execute("SET lock_configuration = true")
total_column = (
f"__govoplan_internal_total_rows_{secrets.token_hex(16)}"
)
query = (
f'SELECT *, COUNT(*) OVER () AS "{total_column}" '
f'FROM ({payload["sql"]}) AS "_govoplan_result" '
f'LIMIT {int(payload["max_rows"]) + 1}'
)
table = connection.execute(query).to_arrow_table()
total_rows = (
int(table.column(total_column)[0].as_py())
if table.num_rows
else 0
)
table = table.drop([total_column])
truncated = table.num_rows > int(payload["max_rows"])
if truncated:
table = table.slice(0, int(payload["max_rows"]))
if table.nbytes > int(payload["max_bytes"]):
raise MemoryError(
"DuckDB output exceeds the configured byte budget."
)
sink = pa.BufferOutputStream()
with pa.ipc.new_stream(sink, table.schema) as writer:
writer.write_table(table)
return {
"ipc": sink.getvalue().to_pybytes(),
"total_rows": total_rows,
"truncated": truncated,
"duckdb_version": duckdb.__version__,
}
finally:
connection.close()
__all__ = [
"DUCKDB_BACKEND_VERSION",
"DuckDbExecutionBackend",
]
@@ -0,0 +1,156 @@
from __future__ import annotations
from govoplan_dataflow.backend.backends.base import (
BackendExecutionError,
BackendExecutionRequest,
BackendExecutionResult,
BackendSource,
canonical_result_schema,
)
from govoplan_dataflow.backend.batches import TypedBatch
from govoplan_dataflow.backend.executor import (
EXECUTOR_VERSION,
PipelineExecutionError,
ResolvedSource,
execute_preview,
)
from govoplan_dataflow.backend.ir import IrExecutionResult, ir_to_graph
from govoplan_dataflow.backend.schemas import GraphNode
class ReferenceExecutionBackend:
name = "reference"
version = EXECUTOR_VERSION
def available(self) -> bool:
return True
def supports(self, request: BackendExecutionRequest) -> bool:
return request.graph.ir_version == 1
def execute(
self,
request: BackendExecutionRequest,
) -> BackendExecutionResult:
_validate_source_batches(request)
try:
result = execute_preview(
ir_to_graph(request.graph),
row_limit=request.budget.max_output_rows,
source_resolver=lambda node, limit: _resolve_source(
request,
node,
limit,
),
preview_node_id=request.preview_node_id,
)
except PipelineExecutionError as exc:
raise BackendExecutionError(
str(exc),
code="backend.reference",
node_id=exc.node_id,
diagnostics=tuple(exc.diagnostics),
) from exc
observed_batch = TypedBatch.from_rows(result.rows)
batch = TypedBatch.from_rows(
result.rows,
schema=canonical_result_schema(
request,
observed_batch.schema,
),
)
try:
batch.ensure_within(
max_rows=request.budget.max_output_rows,
max_bytes=request.budget.max_batch_bytes,
)
except ValueError as exc:
raise BackendExecutionError(
str(exc),
code="backend.budget",
) from exc
lineage = request.graph.lineage.model_copy(
update={
"source_fingerprints": tuple(result.source_fingerprints)
}
)
contract = IrExecutionResult(
backend=self.name,
backend_version=self.version,
graph_semantic_hash=request.graph.semantic_hash,
result_schema=batch.schema,
row_count=result.total_rows,
byte_count=batch.byte_count,
truncated=result.truncated,
diagnostics=tuple(result.diagnostics),
lineage=lineage,
)
return BackendExecutionResult(
contract=contract,
batch=batch,
node_diagnostics=tuple(result.node_diagnostics),
node_preview=result.node_preview,
metadata={
"input_row_count": result.input_row_count,
"source_fingerprints": tuple(result.source_fingerprints),
},
)
def _validate_source_batches(request: BackendExecutionRequest) -> None:
for source in request.sources.values():
try:
source.batch.ensure_within(
max_rows=10_000,
max_bytes=request.budget.max_batch_bytes,
)
except ValueError as exc:
raise BackendExecutionError(
str(exc),
code="backend.source_budget",
node_id=source.node_id,
) from exc
def _resolve_source(
request: BackendExecutionRequest,
node: GraphNode,
limit: int,
) -> ResolvedSource:
source = _source_for_node(request, node)
if source is None:
raise PipelineExecutionError(
f"No typed source batch was supplied for node {node.id!r}.",
node_id=node.id,
)
rows = source.batch.to_rows()
return ResolvedSource(
rows=tuple(rows[:limit]),
source_ref=source.source_ref,
provider=source.provider,
fingerprint=source.fingerprint,
total_rows=source.total_rows,
truncated=source.truncated or len(rows) > limit,
)
def _source_for_node(
request: BackendExecutionRequest,
node: GraphNode,
) -> BackendSource | None:
candidates = (
node.id,
str(node.config.get("source_ref") or ""),
str(node.config.get("source_name") or ""),
)
return next(
(
request.sources[key]
for key in candidates
if key and key in request.sources
),
None,
)
__all__ = ["ReferenceExecutionBackend"]
@@ -0,0 +1,157 @@
from __future__ import annotations
import threading
from typing import Mapping
from govoplan_dataflow.backend.backends.base import (
BackendExecutionError,
BackendExecutionRequest,
BackendExecutionResult,
BackendSource,
ExecutionBackend,
ExecutionBudget,
)
from govoplan_dataflow.backend.backends.duckdb import DuckDbExecutionBackend
from govoplan_dataflow.backend.backends.reference import ReferenceExecutionBackend
from govoplan_dataflow.backend.ir import TypedGraphIr, graph_to_ir
from govoplan_dataflow.backend.planner import PlanningError, plan_execution
from govoplan_dataflow.backend.schemas import PipelineGraph
class ExecutionBackendRegistry:
def __init__(self) -> None:
self._backends: dict[str, ExecutionBackend] = {}
self._slots: dict[str, threading.BoundedSemaphore] = {}
def register(
self,
backend: ExecutionBackend,
*,
max_concurrency: int,
) -> None:
if backend.name in self._backends:
raise ValueError(
f"Duplicate Dataflow execution backend {backend.name!r}."
)
if max_concurrency < 1:
raise ValueError("Backend concurrency must be positive.")
self._backends[backend.name] = backend
self._slots[backend.name] = threading.BoundedSemaphore(
max_concurrency
)
def names(self) -> tuple[str, ...]:
return tuple(self._backends)
def available(self) -> tuple[str, ...]:
return tuple(
name
for name, backend in self._backends.items()
if backend.available()
)
def resolve(
self,
name: str,
request: BackendExecutionRequest,
) -> ExecutionBackend:
if name == "auto":
return self._automatic_backend(request)
backend = self._backends.get(name)
if backend is None:
raise BackendExecutionError(
f"Unknown Dataflow execution backend {name!r}.",
code="backend.unknown",
)
if not backend.available():
raise BackendExecutionError(
f"Dataflow execution backend {name!r} is unavailable.",
code="backend.unavailable",
)
if not backend.supports(request):
raise BackendExecutionError(
f"Dataflow execution backend {name!r} does not support this graph.",
code="backend.unsupported_graph",
)
return backend
def execute(
self,
name: str,
request: BackendExecutionRequest,
) -> BackendExecutionResult:
backend = self.resolve(name, request)
slot = self._slots[backend.name]
acquired = slot.acquire(timeout=request.budget.max_wall_seconds)
if not acquired:
raise BackendExecutionError(
f"Dataflow execution backend {backend.name!r} is at capacity.",
code="backend.capacity",
)
try:
return backend.execute(request)
finally:
slot.release()
def _automatic_backend(
self,
request: BackendExecutionRequest,
) -> ExecutionBackend:
for name in ("duckdb", "reference"):
backend = self._backends.get(name)
if (
backend is not None
and backend.available()
and backend.supports(request)
):
return backend
raise BackendExecutionError(
"No Dataflow execution backend supports this graph.",
code="backend.unavailable",
)
EXECUTION_BACKENDS = ExecutionBackendRegistry()
EXECUTION_BACKENDS.register(
ReferenceExecutionBackend(),
max_concurrency=4,
)
EXECUTION_BACKENDS.register(
DuckDbExecutionBackend(),
max_concurrency=1,
)
def execute_typed_graph(
graph: PipelineGraph | TypedGraphIr,
*,
backend: str = "auto",
sources: Mapping[str, BackendSource] | None = None,
budget: ExecutionBudget | None = None,
preview_node_id: str | None = None,
registry: ExecutionBackendRegistry = EXECUTION_BACKENDS,
) -> BackendExecutionResult:
typed_graph = graph if isinstance(graph, TypedGraphIr) else graph_to_ir(graph)
resolved_budget = budget or ExecutionBudget()
try:
plan = plan_execution(typed_graph)
except PlanningError as exc:
raise BackendExecutionError(
str(exc),
code="backend.planning",
diagnostics=exc.diagnostics,
) from exc
request = BackendExecutionRequest(
plan=plan,
sources=sources or {},
budget=resolved_budget,
preview_node_id=preview_node_id,
)
return registry.execute(backend, request)
__all__ = [
"EXECUTION_BACKENDS",
"ExecutionBackendRegistry",
"execute_typed_graph",
]
+222
View File
@@ -0,0 +1,222 @@
from __future__ import annotations
import json
from dataclasses import dataclass
from typing import Any, Iterable, Mapping
from govoplan_dataflow.backend.ir import (
IrField,
IrSchema,
schema_from_fields,
schema_from_rows,
)
class ArrowDependencyError(RuntimeError):
pass
@dataclass(frozen=True, slots=True)
class TypedBatch:
schema: IrSchema
columns: Mapping[str, tuple[Any, ...]]
row_count: int
byte_count: int
@classmethod
def from_rows(
cls,
rows: Iterable[Mapping[str, Any]],
*,
schema: IrSchema | None = None,
) -> "TypedBatch":
normalized = [dict(row) for row in rows]
resolved_schema = schema or schema_from_rows(normalized)
names = _column_names(normalized, resolved_schema)
columns = {
name: tuple(row.get(name) for row in normalized)
for name in names
}
return cls(
schema=resolved_schema,
columns=columns,
row_count=len(normalized),
byte_count=_estimated_bytes(normalized),
)
@classmethod
def from_arrow_ipc(cls, payload: bytes) -> "TypedBatch":
pa = _pyarrow()
with pa.ipc.open_stream(payload) as reader:
table = reader.read_all()
return cls.from_arrow_table(table)
@classmethod
def from_arrow_table(cls, table: object) -> "TypedBatch":
pa = _pyarrow()
if not isinstance(table, pa.Table):
raise TypeError("Expected a pyarrow.Table.")
schema = schema_from_fields(
tuple(
IrField(
name=field.name,
type=_arrow_data_type(field.type),
nullable=field.nullable,
)
for field in table.schema
),
)
columns = {
name: tuple(table.column(name).to_pylist())
for name in table.column_names
}
return cls(
schema=schema,
columns=columns,
row_count=table.num_rows,
byte_count=table.nbytes,
)
def to_rows(self) -> list[dict[str, Any]]:
names = list(self.columns)
return [
{
name: self.columns[name][index]
for name in names
}
for index in range(self.row_count)
]
def to_arrow_table(self) -> object:
pa = _pyarrow()
if self.columns:
return pa.table(
{
name: list(values)
for name, values in self.columns.items()
}
)
fields = [
pa.field(
field.name,
_pyarrow_type(field.type),
nullable=field.nullable,
)
for field in self.schema.fields
]
schema = pa.schema(fields)
return pa.Table.from_arrays(
[pa.array([], type=field.type) for field in fields],
schema=schema,
)
def to_arrow_ipc(self) -> bytes:
pa = _pyarrow()
table = self.to_arrow_table()
sink = pa.BufferOutputStream()
with pa.ipc.new_stream(sink, table.schema) as writer:
writer.write_table(table)
return sink.getvalue().to_pybytes()
def ensure_within(
self,
*,
max_rows: int,
max_bytes: int,
) -> None:
if self.row_count > max_rows:
raise ValueError(
f"Typed batch has {self.row_count} rows; limit is {max_rows}."
)
if self.byte_count > max_bytes:
raise ValueError(
f"Typed batch uses {self.byte_count} bytes; limit is {max_bytes}."
)
def arrow_available() -> bool:
try:
_pyarrow()
except ArrowDependencyError:
return False
return True
def _column_names(
rows: list[dict[str, Any]],
schema: IrSchema,
) -> list[str]:
names = [field.name for field in schema.fields]
seen = set(names)
for row in rows:
for name in row:
if name not in seen:
names.append(name)
seen.add(name)
return names
def _estimated_bytes(rows: list[dict[str, Any]]) -> int:
encoded = json.dumps(
rows,
sort_keys=True,
separators=(",", ":"),
ensure_ascii=True,
default=str,
)
return len(encoded.encode("utf-8"))
def _pyarrow() -> Any:
try:
import pyarrow as pa
except ImportError as exc:
raise ArrowDependencyError(
"Arrow execution requires the govoplan-dataflow analytics extra."
) from exc
return pa
def _arrow_data_type(data_type: object) -> str:
pa = _pyarrow()
if pa.types.is_boolean(data_type):
return "boolean"
if pa.types.is_integer(data_type):
return "integer"
if pa.types.is_floating(data_type) or pa.types.is_decimal(data_type):
return "number"
if pa.types.is_date(data_type):
return "date"
if pa.types.is_timestamp(data_type):
return "datetime"
if pa.types.is_binary(data_type):
return "binary"
if pa.types.is_string(data_type) or pa.types.is_large_string(data_type):
return "string"
if pa.types.is_null(data_type):
return "null"
return "unknown"
def _pyarrow_type(data_type: str) -> object:
pa = _pyarrow()
mapping = {
"boolean": pa.bool_(),
"integer": pa.int64(),
"number": pa.float64(),
"date": pa.date32(),
"datetime": pa.timestamp("us", tz="UTC"),
"binary": pa.binary(),
"string": pa.string(),
"null": pa.null(),
"json": pa.string(),
"unknown": pa.null(),
}
return mapping.get(data_type, pa.null())
__all__ = [
"ArrowDependencyError",
"TypedBatch",
"arrow_available",
]
+491
View File
@@ -0,0 +1,491 @@
from __future__ import annotations
import hashlib
import json
from typing import Any, Literal
from pydantic import BaseModel, ConfigDict, Field
from govoplan_dataflow.backend.expressions import (
ExpressionError,
infer_expression_type,
parse_expression,
)
from govoplan_dataflow.backend.graph import topological_order
from govoplan_dataflow.backend.node_library import node_definition
from govoplan_dataflow.backend.schema_validation import (
SchemaState,
propagate_graph_schemas,
)
from govoplan_dataflow.backend.schemas import (
DataflowDiagnostic,
GraphEdge,
GraphNode,
GraphPosition,
PipelineGraph,
)
DATAFLOW_IR_VERSION = 1
OPERATOR_CONTRACT_VERSION = 1
DataType = Literal[
"unknown",
"null",
"string",
"integer",
"number",
"boolean",
"date",
"datetime",
"binary",
"json",
]
class IrModel(BaseModel):
model_config = ConfigDict(frozen=True)
class IrField(IrModel):
name: str = Field(min_length=1, max_length=300)
type: DataType = "unknown"
nullable: bool = True
metadata: dict[str, Any] = Field(default_factory=dict)
class IrSchema(IrModel):
fields: tuple[IrField, ...] = ()
open: bool = False
semantic_hash: str = Field(default="", max_length=64)
class IrPort(IrModel):
id: str = Field(min_length=1, max_length=80)
direction: Literal["input", "output"]
required: bool = True
multiple: bool = False
minimum_connections: int = Field(default=1, ge=0)
data_schema: IrSchema | None = None
class IrExpression(IrModel):
source: str = Field(min_length=1, max_length=10_000)
dialect: Literal["duckdb"] = "duckdb"
columns: tuple[str, ...] = ()
result_type: DataType = "unknown"
semantic_hash: str = Field(min_length=64, max_length=64)
class IrParameter(IrModel):
id: str = Field(min_length=1, max_length=120)
type: DataType = "unknown"
required: bool = False
default: Any = None
description: str | None = Field(default=None, max_length=1_000)
class IrNode(IrModel):
id: str = Field(min_length=1, max_length=100)
type: str = Field(min_length=1, max_length=100)
operator_version: int = Field(default=OPERATOR_CONTRACT_VERSION, ge=1)
label: str = Field(min_length=1, max_length=300)
position: GraphPosition
config: dict[str, Any] = Field(default_factory=dict)
input_ports: tuple[IrPort, ...] = ()
output_ports: tuple[IrPort, ...] = ()
expressions: tuple[IrExpression, ...] = ()
output_schema: IrSchema = Field(default_factory=IrSchema)
class IrEdge(IrModel):
id: str = Field(min_length=1, max_length=255)
source: str = Field(min_length=1, max_length=100)
target: str = Field(min_length=1, max_length=100)
source_port: str = Field(default="output", min_length=1, max_length=80)
target_port: str = Field(default="input", min_length=1, max_length=80)
class IrLineage(IrModel):
generated_from: Literal["graph", "sql", "template"] = "graph"
source_fingerprints: tuple[dict[str, Any], ...] = ()
parent_semantic_hashes: tuple[str, ...] = ()
loss_diagnostics: tuple[DataflowDiagnostic, ...] = ()
class TypedGraphIr(IrModel):
ir_version: Literal[1] = DATAFLOW_IR_VERSION
graph_schema_version: int = Field(default=1, ge=1)
nodes: tuple[IrNode, ...] = Field(default=(), max_length=100)
edges: tuple[IrEdge, ...] = Field(default=(), max_length=200)
parameters: tuple[IrParameter, ...] = Field(default=(), max_length=100)
lineage: IrLineage = Field(default_factory=IrLineage)
semantic_hash: str = Field(min_length=64, max_length=64)
class IrExecutionResult(IrModel):
ir_version: Literal[1] = DATAFLOW_IR_VERSION
backend: str = Field(min_length=1, max_length=100)
backend_version: str = Field(min_length=1, max_length=100)
graph_semantic_hash: str = Field(min_length=64, max_length=64)
result_schema: IrSchema
row_count: int = Field(ge=0)
byte_count: int = Field(ge=0)
truncated: bool = False
diagnostics: tuple[DataflowDiagnostic, ...] = ()
lineage: IrLineage = Field(default_factory=IrLineage)
def graph_to_ir(
graph: PipelineGraph,
*,
parameters: tuple[IrParameter, ...] = (),
lineage: IrLineage | None = None,
) -> TypedGraphIr:
ordered, cyclic = topological_order(graph)
schemas, schema_diagnostics = (
({}, [])
if cyclic
else propagate_graph_schemas(graph, ordered=ordered)
)
resolved_lineage = _lineage_with_diagnostics(
lineage or IrLineage(),
schema_diagnostics,
)
nodes = tuple(
_node_ir(node, schemas.get(node.id, _unknown_schema()))
for node in graph.nodes
)
edges = tuple(_edge_ir(edge) for edge in graph.edges)
semantic_hash = _graph_semantic_hash(
graph_schema_version=graph.schema_version,
nodes=nodes,
edges=edges,
parameters=parameters,
)
return TypedGraphIr(
graph_schema_version=graph.schema_version,
nodes=nodes,
edges=edges,
parameters=parameters,
lineage=resolved_lineage,
semantic_hash=semantic_hash,
)
def ir_to_graph(ir: TypedGraphIr) -> PipelineGraph:
return PipelineGraph(
schema_version=ir.graph_schema_version,
nodes=[
GraphNode(
id=node.id,
type=node.type,
label=node.label,
position=node.position.model_copy(deep=True),
config=dict(node.config),
)
for node in ir.nodes
],
edges=[
GraphEdge(
id=edge.id,
source=edge.source,
target=edge.target,
source_port=edge.source_port,
target_port=edge.target_port,
)
for edge in ir.edges
],
)
def schema_from_rows(rows: list[dict[str, Any]]) -> IrSchema:
names: list[str] = []
seen: set[str] = set()
for row in rows:
for name in row:
normalized = str(name)
if normalized not in seen:
names.append(normalized)
seen.add(normalized)
columns = {
name: [row.get(name) for row in rows]
for name in names
}
fields = tuple(
IrField(
name=name,
type=_values_type(values),
nullable=any(value is None for value in values),
)
for name, values in columns.items()
)
return _schema(fields)
def schema_from_fields(
fields: tuple[IrField, ...],
*,
open: bool = False,
) -> IrSchema:
return _schema(fields, open=open)
def _node_ir(node: GraphNode, output_state: SchemaState) -> IrNode:
definition = node_definition(node.type)
input_ports = (
tuple(
IrPort(
id=port.id,
direction="input",
required=port.required,
multiple=port.multiple,
minimum_connections=port.minimum_connections,
)
for port in definition.input_ports
)
if definition
else ()
)
output_schema = _state_schema(output_state)
output_ports = (
tuple(
IrPort(
id=port.id,
direction="output",
required=port.required,
multiple=port.multiple,
minimum_connections=port.minimum_connections,
data_schema=output_schema,
)
for port in definition.output_ports
)
if definition
else ()
)
return IrNode(
id=node.id,
type=node.type,
label=node.label,
position=node.position.model_copy(deep=True),
config=dict(node.config),
input_ports=input_ports,
output_ports=output_ports,
expressions=_node_expressions(node, output_state),
output_schema=output_schema,
)
def _node_expressions(
node: GraphNode,
output_state: SchemaState,
) -> tuple[IrExpression, ...]:
source = node.config.get("expression")
if node.type not in {"expression", "filter.expression"} or not source:
return ()
try:
parsed = parse_expression(str(source))
except ExpressionError:
return ()
inferred = infer_expression_type(parsed, output_state.types)
return (
IrExpression(
source=parsed.source,
columns=parsed.columns,
result_type=_data_type(inferred),
semantic_hash=_hash_payload(
{
"dialect": "duckdb",
"source": parsed.sql(),
"result_type": inferred,
}
),
),
)
def _edge_ir(edge: GraphEdge) -> IrEdge:
return IrEdge(
id=edge.id,
source=edge.source,
target=edge.target,
source_port=edge.source_port,
target_port=edge.target_port,
)
def _state_schema(state: SchemaState) -> IrSchema:
return _schema(
tuple(
IrField(
name=name,
type=_data_type(state.type_of(name)),
nullable=True,
)
for name in sorted(state.columns)
),
open=state.open,
)
def _schema(
fields: tuple[IrField, ...],
*,
open: bool = False,
) -> IrSchema:
semantic_hash = _hash_payload(
{
"fields": [
{
"name": item.name,
"type": item.type,
"nullable": item.nullable,
}
for item in fields
],
"open": open,
}
)
return IrSchema(
fields=fields,
open=open,
semantic_hash=semantic_hash,
)
def _graph_semantic_hash(
*,
graph_schema_version: int,
nodes: tuple[IrNode, ...],
edges: tuple[IrEdge, ...],
parameters: tuple[IrParameter, ...],
) -> str:
return _hash_payload(
{
"ir_version": DATAFLOW_IR_VERSION,
"graph_schema_version": graph_schema_version,
"nodes": [
{
"id": node.id,
"type": node.type,
"operator_version": node.operator_version,
"config": node.config,
"expressions": [
expression.model_dump(mode="json")
for expression in node.expressions
],
}
for node in nodes
],
"edges": [
edge.model_dump(mode="json")
for edge in edges
],
"parameters": [
parameter.model_dump(mode="json")
for parameter in parameters
],
}
)
def _lineage_with_diagnostics(
lineage: IrLineage,
diagnostics: list[DataflowDiagnostic],
) -> IrLineage:
if not diagnostics:
return lineage
return lineage.model_copy(
update={
"loss_diagnostics": (
*lineage.loss_diagnostics,
*diagnostics,
)
}
)
def _values_type(values: list[Any]) -> DataType:
concrete = {
_value_type(value)
for value in values
if value is not None
}
if not concrete:
return "null"
if concrete <= {"integer", "number"}:
return "number" if "number" in concrete else "integer"
return concrete.pop() if len(concrete) == 1 else "unknown"
def _value_type(value: Any) -> DataType:
from datetime import date, datetime
from decimal import Decimal
if isinstance(value, bool):
return "boolean"
if isinstance(value, int):
return "integer"
if isinstance(value, (float, Decimal)):
return "number"
if isinstance(value, datetime):
return "datetime"
if isinstance(value, date):
return "date"
if isinstance(value, bytes):
return "binary"
if isinstance(value, str):
return "string"
if isinstance(value, (dict, list, tuple)):
return "json"
return "unknown"
def _data_type(value: str) -> DataType:
known = {
"unknown",
"null",
"string",
"integer",
"number",
"boolean",
"date",
"datetime",
"binary",
"json",
}
return value if value in known else "unknown" # type: ignore[return-value]
def _unknown_schema() -> SchemaState:
return SchemaState(frozenset(), open=True)
def _hash_payload(payload: object) -> str:
encoded = json.dumps(
payload,
sort_keys=True,
separators=(",", ":"),
ensure_ascii=True,
default=str,
)
return hashlib.sha256(encoded.encode("utf-8")).hexdigest()
__all__ = [
"DATAFLOW_IR_VERSION",
"DataType",
"IrEdge",
"IrExecutionResult",
"IrExpression",
"IrField",
"IrLineage",
"IrNode",
"IrParameter",
"IrPort",
"IrSchema",
"TypedGraphIr",
"graph_to_ir",
"ir_to_graph",
"schema_from_fields",
"schema_from_rows",
]
+99
View File
@@ -0,0 +1,99 @@
from __future__ import annotations
import hashlib
from dataclasses import dataclass
from govoplan_dataflow.backend.graph import topological_order, validate_graph
from govoplan_dataflow.backend.ir import TypedGraphIr, ir_to_graph
from govoplan_dataflow.backend.schemas import DataflowDiagnostic
from govoplan_dataflow.backend.sql_compiler import (
SqlCompilationError,
render_sql,
)
@dataclass(frozen=True, slots=True)
class ExecutionPlan:
graph: TypedGraphIr
ordered_node_ids: tuple[str, ...]
diagnostics: tuple[DataflowDiagnostic, ...]
generated_sql: str | None
sql_diagnostics: tuple[DataflowDiagnostic, ...]
semantic_hash: str
class PlanningError(ValueError):
def __init__(self, diagnostics: tuple[DataflowDiagnostic, ...]) -> None:
super().__init__(
diagnostics[0].message
if diagnostics
else "Dataflow planning failed."
)
self.diagnostics = diagnostics
def plan_execution(graph: TypedGraphIr) -> ExecutionPlan:
pipeline_graph = ir_to_graph(graph)
diagnostics = tuple(validate_graph(pipeline_graph))
errors = tuple(
item
for item in diagnostics
if item.severity == "error"
)
if errors:
raise PlanningError(errors)
ordered, cyclic = topological_order(pipeline_graph)
if cyclic:
cycle = DataflowDiagnostic(
severity="error",
code="graph.cycle",
message="A cyclic graph cannot be planned.",
)
raise PlanningError((cycle,))
sql_text, sql_diagnostics = _generated_sql(pipeline_graph)
semantic_hash = _plan_hash(
graph.semantic_hash,
ordered,
sql_text=sql_text,
)
return ExecutionPlan(
graph=graph,
ordered_node_ids=tuple(ordered),
diagnostics=diagnostics,
generated_sql=sql_text,
sql_diagnostics=sql_diagnostics,
semantic_hash=semantic_hash,
)
def _generated_sql(
graph,
) -> tuple[str | None, tuple[DataflowDiagnostic, ...]]:
try:
sql_text, diagnostics = render_sql(graph)
except SqlCompilationError as exc:
return None, tuple(exc.diagnostics)
return sql_text, tuple(diagnostics)
def _plan_hash(
graph_hash: str,
ordered: list[str],
*,
sql_text: str | None,
) -> str:
payload = "\n".join(
(
graph_hash,
*ordered,
(sql_text or "").strip(),
)
)
return hashlib.sha256(payload.encode("utf-8")).hexdigest()
__all__ = [
"ExecutionPlan",
"PlanningError",
"plan_execution",
]
@@ -55,6 +55,15 @@ def validate_graph_schemas(
*, *,
ordered: list[str], ordered: list[str],
) -> list[DataflowDiagnostic]: ) -> 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} node_by_id = {node.id: node for node in graph.nodes}
inputs = _graph_inputs_by_port(graph) inputs = _graph_inputs_by_port(graph)
schemas: dict[str, SchemaState] = {} schemas: dict[str, SchemaState] = {}
@@ -81,7 +90,7 @@ def validate_graph_schemas(
) )
schemas[node.id] = result.state schemas[node.id] = result.state
diagnostics.extend(result.diagnostics) diagnostics.extend(result.diagnostics)
return diagnostics return schemas, diagnostics
def _graph_inputs_by_port( def _graph_inputs_by_port(
+1
View File
@@ -193,6 +193,7 @@ class PipelinePreviewRequest(BaseModel):
pattern=r"^[A-Za-z0-9_.:-]+$", pattern=r"^[A-Za-z0-9_.:-]+$",
) )
row_limit: int = Field(default=100, ge=1, le=500) row_limit: int = Field(default=100, ge=1, le=500)
execution_backend: Literal["reference", "duckdb", "auto"] = "reference"
class PreviewColumn(BaseModel): class PreviewColumn(BaseModel):
+177 -43
View File
@@ -25,6 +25,13 @@ from govoplan_core.core.datasources import (
datasource_publication, datasource_publication,
) )
from govoplan_core.db.base import utcnow from govoplan_core.db.base import utcnow
from govoplan_dataflow.backend.backends import (
BackendExecutionError,
BackendSource,
ExecutionBudget,
execute_typed_graph,
)
from govoplan_dataflow.backend.batches import TypedBatch
from govoplan_dataflow.backend.db.models import ( from govoplan_dataflow.backend.db.models import (
DataflowPipeline, DataflowPipeline,
DataflowPipelineRevision, DataflowPipelineRevision,
@@ -32,6 +39,7 @@ from govoplan_dataflow.backend.db.models import (
) )
from govoplan_dataflow.backend.executor import ( from govoplan_dataflow.backend.executor import (
EXECUTOR_VERSION, EXECUTOR_VERSION,
MAX_SOURCE_ROWS,
PipelineExecutionError, PipelineExecutionError,
PipelineExecutionResult, PipelineExecutionResult,
ResolvedSource, ResolvedSource,
@@ -50,6 +58,7 @@ from govoplan_dataflow.backend.graph import (
from govoplan_dataflow.backend.schemas import ( from govoplan_dataflow.backend.schemas import (
DataflowDiagnostic, DataflowDiagnostic,
GraphNode, GraphNode,
NodePreviewResult,
PipelineCreateRequest, PipelineCreateRequest,
PipelineDeriveRequest, PipelineDeriveRequest,
PipelineDraftRequest, PipelineDraftRequest,
@@ -62,6 +71,7 @@ from govoplan_dataflow.backend.schemas import (
PipelineSqlResponse, PipelineSqlResponse,
PipelineUpdateRequest, PipelineUpdateRequest,
PipelineValidationResponse, PipelineValidationResponse,
PreviewColumn,
) )
from govoplan_dataflow.backend.sql_compiler import ( from govoplan_dataflow.backend.sql_compiler import (
SqlCompilationError, SqlCompilationError,
@@ -588,49 +598,13 @@ def preview_pipeline(
started_at = utcnow() started_at = utcnow()
run: DataflowRun | None = None run: DataflowRun | None = None
try: try:
provider = datasource_catalogue(registry) result, executor_version = _execute_pipeline_preview(
def resolve_source(node: GraphNode, limit: int) -> ResolvedSource:
if provider is None:
raise PipelineExecutionError(
"Datasource-backed preview requires the Datasources catalogue capability.",
node_id=node.id,
)
if principal is None:
raise PipelineExecutionError(
"Datasource-backed preview requires a tenant API principal.",
node_id=node.id,
)
try:
resolved = provider.read_datasource(
session,
principal,
request=DatasourceReadRequest(
datasource_ref=str(node.config["source_ref"]),
consistency=str(
node.config.get("consistency") or "current"
), # type: ignore[arg-type]
limit=limit,
expected_fingerprint=_clean_optional(
node.config.get("expected_fingerprint")
),
),
)
except DatasourceError as exc:
raise PipelineExecutionError(str(exc), node_id=node.id) from exc
return ResolvedSource(
rows=tuple(dict(row) for row in resolved.rows),
source_ref=resolved.datasource.ref,
provider=resolved.datasource.provider or "datasources",
fingerprint=resolved.datasource.fingerprint,
total_rows=resolved.total_rows,
truncated=resolved.truncated,
)
result = execute_preview(
graph, graph,
session=session,
principal=principal,
registry=registry,
backend=payload.execution_backend,
row_limit=payload.row_limit, row_limit=payload.row_limit,
source_resolver=resolve_source,
preview_node_id=payload.preview_node_id, preview_node_id=payload.preview_node_id,
) )
status = "succeeded" status = "succeeded"
@@ -645,6 +619,11 @@ def preview_pipeline(
source_fingerprints = result.source_fingerprints source_fingerprints = result.source_fingerprints
input_row_count = result.input_row_count input_row_count = result.input_row_count
except PipelineExecutionError as exc: except PipelineExecutionError as exc:
executor_version = (
payload.execution_backend
if payload.execution_backend != "reference"
else EXECUTOR_VERSION
)
status = "failed" status = "failed"
error = str(exc) error = str(exc)
diagnostics = [ diagnostics = [
@@ -672,7 +651,7 @@ def preview_pipeline(
pipeline_revision_id=revision.id, pipeline_revision_id=revision.id,
run_type="preview", run_type="preview",
status=status, status=status,
executor_version=EXECUTOR_VERSION, executor_version=executor_version,
definition_hash=graph_hash, definition_hash=graph_hash,
source_fingerprints=source_fingerprints, source_fingerprints=source_fingerprints,
result_schema=[item.model_dump(mode="json") for item in columns], result_schema=[item.model_dump(mode="json") for item in columns],
@@ -702,7 +681,162 @@ def preview_pipeline(
source_fingerprints=source_fingerprints, source_fingerprints=source_fingerprints,
input_row_count=input_row_count, input_row_count=input_row_count,
definition_hash=graph_hash, definition_hash=graph_hash,
executor_version=EXECUTOR_VERSION, executor_version=executor_version,
)
def _execute_pipeline_preview(
graph: PipelineGraph,
*,
session: Session,
principal: ApiPrincipal | None,
registry: object | None,
backend: str,
row_limit: int,
preview_node_id: str | None,
) -> tuple[PipelineExecutionResult, str]:
source_resolver = _preview_source_resolver(
session=session,
principal=principal,
registry=registry,
)
if backend == "reference":
return (
execute_preview(
graph,
row_limit=row_limit,
source_resolver=source_resolver,
preview_node_id=preview_node_id,
),
EXECUTOR_VERSION,
)
sources = _typed_backend_sources(
graph,
source_resolver=source_resolver,
)
try:
result = execute_typed_graph(
graph,
backend=backend,
sources=sources,
budget=ExecutionBudget(max_output_rows=row_limit),
preview_node_id=preview_node_id,
)
except BackendExecutionError as exc:
raise PipelineExecutionError(
str(exc),
node_id=exc.node_id,
diagnostics=tuple(exc.diagnostics),
) from exc
columns = [
PreviewColumn(
name=field.name,
type=field.type,
nullable=field.nullable,
)
for field in result.batch.schema.fields
]
source_fingerprints = list(
result.metadata.get(
"source_fingerprints",
result.contract.lineage.source_fingerprints,
)
)
node_preview = _typed_node_preview(
result,
preview_node_id=preview_node_id,
columns=columns,
)
return (
PipelineExecutionResult(
rows=result.rows,
total_rows=result.contract.row_count,
truncated=result.contract.truncated,
columns=columns,
diagnostics=list(result.contract.diagnostics),
node_diagnostics=list(result.node_diagnostics),
node_preview=node_preview,
source_fingerprints=source_fingerprints,
input_row_count=int(
result.metadata.get(
"input_row_count",
sum(
int(
item.get(
"row_count",
item.get("total_rows", 0),
)
)
for item in source_fingerprints
),
)
),
),
result.contract.backend_version,
)
def _preview_source_resolver(
*,
session: Session,
principal: ApiPrincipal | None,
registry: object | None,
):
if principal is not None:
return _datasource_source_resolver(
session=session,
principal=principal,
registry=registry,
)
def unavailable(node: GraphNode, _limit: int) -> ResolvedSource:
raise PipelineExecutionError(
"Datasource-backed preview requires a tenant API principal.",
node_id=node.id,
)
return unavailable
def _typed_backend_sources(
graph: PipelineGraph,
*,
source_resolver,
) -> dict[str, BackendSource]:
sources: dict[str, BackendSource] = {}
for node in graph.nodes:
if node.type != "source.reference":
continue
resolved = source_resolver(node, MAX_SOURCE_ROWS)
sources[node.id] = BackendSource(
node_id=node.id,
batch=TypedBatch.from_rows(resolved.rows),
source_ref=resolved.source_ref,
provider=resolved.provider,
fingerprint=resolved.fingerprint,
total_rows=resolved.total_rows,
truncated=resolved.truncated,
source_name=str(node.config.get("source_name") or ""),
)
return sources
def _typed_node_preview(
result,
*,
preview_node_id: str | None,
columns: list[PreviewColumn],
) -> NodePreviewResult | None:
if result.node_preview is not None:
return result.node_preview
if preview_node_id is None:
return None
return NodePreviewResult(
node_id=preview_node_id,
columns=columns,
rows=result.rows,
total_rows=result.contract.row_count,
truncated=result.contract.truncated,
) )
@@ -1409,7 +1409,7 @@ def _render_select(
else str(field_config.get("alias") or column) else str(field_config.get("alias") or column)
) )
expression: exp.Expression = column_expression(column) expression: exp.Expression = column_expression(column)
if alias != column: if alias != expression.alias_or_name:
expression = expression.as_(alias) expression = expression.as_(alias)
state.select_expressions.append(expression) state.select_expressions.append(expression)
state.selected = True state.selected = True
+32
View File
@@ -18,6 +18,7 @@ from govoplan_core.core.datasources import (
DatasourcePublicationResult, DatasourcePublicationResult,
) )
from govoplan_core.db.base import Base from govoplan_core.db.base import Base
from govoplan_dataflow.backend.backends.duckdb import DuckDbExecutionBackend
from govoplan_dataflow.backend.db.models import ( from govoplan_dataflow.backend.db.models import (
DataflowPipeline, DataflowPipeline,
DataflowPipelineRevision, DataflowPipelineRevision,
@@ -320,6 +321,37 @@ class DataflowServiceTests(unittest.TestCase):
self.session.scalar(select(func.count()).select_from(DataflowRun)), self.session.scalar(select(func.count()).select_from(DataflowRun)),
) )
@unittest.skipUnless(
DuckDbExecutionBackend().available(),
"analytics extra is not installed",
)
def test_saved_preview_can_use_isolated_duckdb_backend(self) -> None:
pipeline = self._create()
response = preview_pipeline(
self.session,
tenant_id="tenant-1",
actor_id="user-1",
payload=PipelinePreviewRequest(
pipeline_id=pipeline.id,
execution_backend="duckdb",
),
principal=principal(),
)
self.session.commit()
run = self.session.get(DataflowRun, response.run_id)
self.assertEqual("succeeded", response.status)
self.assertEqual(
[{"id": 2, "amount": 15}, {"id": 3, "amount": 25}],
response.rows,
)
self.assertEqual(2, response.total_rows)
self.assertFalse(response.truncated)
self.assertEqual(3, response.input_row_count)
self.assertTrue(response.executor_version.startswith("duckdb-isolated-v1"))
self.assertEqual(response.executor_version, run.executor_version)
def test_failed_preview_keeps_upstream_lineage_and_failed_node_diagnostic(self) -> None: def test_failed_preview_keeps_upstream_lineage_and_failed_node_diagnostic(self) -> None:
graph = sample_graph() graph = sample_graph()
graph.nodes[0].config["rows"] = [{"id": 1, "amount": "not-a-number"}] graph.nodes[0].config["rows"] = [{"id": 1, "amount": "not-a-number"}]
+258
View File
@@ -0,0 +1,258 @@
from __future__ import annotations
import unittest
from govoplan_dataflow.backend.backends import (
EXECUTION_BACKENDS,
ExecutionBudget,
execute_typed_graph,
)
from govoplan_dataflow.backend.backends.duckdb import DuckDbExecutionBackend
from govoplan_dataflow.backend.batches import TypedBatch
from govoplan_dataflow.backend.ir import (
DATAFLOW_IR_VERSION,
graph_to_ir,
ir_to_graph,
)
from govoplan_dataflow.backend.planner import plan_execution
from govoplan_dataflow.backend.schemas import (
GraphEdge,
GraphNode,
GraphPosition,
PipelineGraph,
)
from govoplan_dataflow.backend.sql_compiler import compile_sql
def inline_source(
node_id: str = "source",
source_name: str = "records",
rows: list[dict] | None = None,
) -> GraphNode:
return GraphNode(
id=node_id,
type="source.inline",
label=source_name,
position=GraphPosition(x=40, y=120),
config={
"source_name": source_name,
"rows": rows
or [
{"department": "a", "amount": 10},
{"department": "a", "amount": 5},
{"department": "b", "amount": 20},
],
},
)
def simple_graph() -> PipelineGraph:
source = inline_source()
derive = GraphNode(
id="derive",
type="expression",
label="Double amount",
position=GraphPosition(x=280, y=120),
config={
"target_column": "double_amount",
"expression": "amount * 2",
"result_type": "integer",
},
)
output = GraphNode(
id="output",
type="output",
label="Output",
position=GraphPosition(x=520, y=120),
config={},
)
return PipelineGraph(
nodes=[source, derive, output],
edges=[
GraphEdge(id="edge-1", source=source.id, target=derive.id),
GraphEdge(id="edge-2", source=derive.id, target=output.id),
],
)
class TypedIrTests(unittest.TestCase):
def test_ir_roundtrip_preserves_identity_layout_and_expressions(self) -> None:
graph = simple_graph()
ir = graph_to_ir(graph)
roundtrip = ir_to_graph(ir)
self.assertEqual(DATAFLOW_IR_VERSION, ir.ir_version)
self.assertEqual(64, len(ir.semantic_hash))
self.assertEqual(
graph.model_dump(mode="json"),
roundtrip.model_dump(mode="json"),
)
expression = ir.nodes[1].expressions[0]
self.assertEqual(("amount",), expression.columns)
self.assertEqual("integer", expression.result_type)
self.assertEqual(
{"amount", "department", "double_amount"},
{field.name for field in ir.nodes[-1].output_schema.fields},
)
def test_semantic_hash_ignores_layout_but_changes_with_behavior(self) -> None:
graph = simple_graph()
moved = graph.model_copy(deep=True)
moved.nodes[1].position = GraphPosition(x=900, y=700)
changed = graph.model_copy(deep=True)
changed.nodes[1].config["expression"] = "amount * 3"
self.assertEqual(
graph_to_ir(graph).semantic_hash,
graph_to_ir(moved).semantic_hash,
)
self.assertNotEqual(
graph_to_ir(graph).semantic_hash,
graph_to_ir(changed).semantic_hash,
)
def test_typed_batch_is_columnar_and_roundtrips_rows(self) -> None:
rows = [
{"name": "Ada", "active": True},
{"name": "Grace", "active": None},
]
batch = TypedBatch.from_rows(rows)
self.assertEqual(2, batch.row_count)
self.assertEqual(("Ada", "Grace"), batch.columns["name"])
self.assertEqual(rows, batch.to_rows())
self.assertGreater(batch.byte_count, 0)
def test_physical_plan_pins_order_sql_and_semantic_hash(self) -> None:
plan = plan_execution(graph_to_ir(simple_graph()))
self.assertEqual(
("source", "derive", "output"),
plan.ordered_node_ids,
)
self.assertIn("amount * 2", plan.generated_sql)
self.assertEqual(64, len(plan.semantic_hash))
self.assertEqual((), plan.sql_diagnostics)
class ExecutionBackendTests(unittest.TestCase):
def test_reference_backend_executes_typed_ir_with_budgets(self) -> None:
result = execute_typed_graph(
simple_graph(),
backend="reference",
budget=ExecutionBudget(max_output_rows=2),
)
self.assertEqual("reference", result.contract.backend)
self.assertEqual(3, result.contract.row_count)
self.assertTrue(result.contract.truncated)
self.assertEqual(
[
{"department": "a", "amount": 10, "double_amount": 20},
{"department": "a", "amount": 5, "double_amount": 10},
],
result.rows,
)
def test_backend_registry_keeps_reference_available(self) -> None:
self.assertIn("reference", EXECUTION_BACKENDS.names())
self.assertIn("reference", EXECUTION_BACKENDS.available())
@unittest.skipUnless(
DuckDbExecutionBackend().available(),
"analytics extra is not installed",
)
def test_duckdb_matches_reference_for_aggregate_join_and_union(self) -> None:
graphs = [
compile_sql(
"""
SELECT department, COUNT(*) AS records, SUM(amount) AS total
FROM records
GROUP BY department
ORDER BY department
""",
source_nodes=[inline_source()],
)[0],
compile_sql(
"""
SELECT records.department, labels.label
FROM records
LEFT JOIN labels ON records.department = labels.department
ORDER BY records.department
""",
source_nodes=[
inline_source(),
inline_source(
"labels-source",
"labels",
[
{"department": "a", "label": "Alpha"},
{"department": "b", "label": "Beta"},
],
),
],
)[0],
compile_sql(
"""
SELECT * FROM first
UNION ALL BY NAME
SELECT * FROM second
ORDER BY amount
""",
source_nodes=[
inline_source(
"first-source",
"first",
[{"amount": 2}, {"amount": None}],
),
inline_source(
"second-source",
"second",
[{"amount": 1}],
),
],
)[0],
]
for graph in graphs:
with self.subTest(nodes=[node.type for node in graph.nodes]):
reference = execute_typed_graph(
graph,
backend="reference",
budget=ExecutionBudget(max_output_rows=100),
)
analytical = execute_typed_graph(
graph,
backend="duckdb",
budget=ExecutionBudget(
max_output_rows=100,
max_wall_seconds=10,
),
)
self.assertEqual(reference.rows, analytical.rows)
self.assertEqual(
reference.contract.row_count,
analytical.contract.row_count,
)
self.assertEqual(
[field.name for field in reference.batch.schema.fields],
[field.name for field in analytical.batch.schema.fields],
)
self.assertEqual(
reference.batch.schema.semantic_hash,
analytical.batch.schema.semantic_hash,
)
self.assertEqual(
reference.contract.diagnostics,
analytical.contract.diagnostics,
)
self.assertEqual(
reference.contract.lineage.source_fingerprints,
analytical.contract.lineage.source_fingerprints,
)
if __name__ == "__main__":
unittest.main()
+1
View File
@@ -478,6 +478,7 @@ export function previewDataflowPipeline(
source_nodes?: PipelineGraphNode[]; source_nodes?: PipelineGraphNode[];
preview_node_id?: string; preview_node_id?: string;
row_limit?: number; row_limit?: number;
execution_backend?: "reference" | "duckdb" | "auto";
} }
): Promise<PipelinePreview> { ): Promise<PipelinePreview> {
return apiFetch<PipelinePreview>(settings, "/api/v1/dataflow/preview", { return apiFetch<PipelinePreview>(settings, "/api/v1/dataflow/preview", {