Add typed Arrow execution backend boundary

2026-07-29 16:47:14 +02:00
parent 69509d5cc2
commit 8a11d7571b
17 changed files with 2264 additions and 49 deletions
@@ -25,6 +25,13 @@ from govoplan_core.core.datasources import (
    datasource_publication,
 )
 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 (
    DataflowPipeline,
    DataflowPipelineRevision,
@@ -32,6 +39,7 @@ from govoplan_dataflow.backend.db.models import (
 )
 from govoplan_dataflow.backend.executor import (
    EXECUTOR_VERSION,
+    MAX_SOURCE_ROWS,
    PipelineExecutionError,
    PipelineExecutionResult,
    ResolvedSource,
@@ -50,6 +58,7 @@ from govoplan_dataflow.backend.graph import (
 from govoplan_dataflow.backend.schemas import (
    DataflowDiagnostic,
    GraphNode,
+    NodePreviewResult,
    PipelineCreateRequest,
    PipelineDeriveRequest,
    PipelineDraftRequest,
@@ -62,6 +71,7 @@ from govoplan_dataflow.backend.schemas import (
    PipelineSqlResponse,
    PipelineUpdateRequest,
    PipelineValidationResponse,
+    PreviewColumn,
 )
 from govoplan_dataflow.backend.sql_compiler import (
    SqlCompilationError,
@@ -588,49 +598,13 @@ def preview_pipeline(
    started_at = utcnow()
    run: DataflowRun | None = None
    try:
-        provider = datasource_catalogue(registry)
-
-        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(
+        result, executor_version = _execute_pipeline_preview(
            graph,
+            session=session,
+            principal=principal,
+            registry=registry,
+            backend=payload.execution_backend,
            row_limit=payload.row_limit,
-            source_resolver=resolve_source,
            preview_node_id=payload.preview_node_id,
        )
        status = "succeeded"
@@ -645,6 +619,11 @@ def preview_pipeline(
        source_fingerprints = result.source_fingerprints
        input_row_count = result.input_row_count
    except PipelineExecutionError as exc:
+        executor_version = (
+            payload.execution_backend
+            if payload.execution_backend != "reference"
+            else EXECUTOR_VERSION
+        )
        status = "failed"
        error = str(exc)
        diagnostics = [
@@ -672,7 +651,7 @@ def preview_pipeline(
            pipeline_revision_id=revision.id,
            run_type="preview",
            status=status,
-            executor_version=EXECUTOR_VERSION,
+            executor_version=executor_version,
            definition_hash=graph_hash,
            source_fingerprints=source_fingerprints,
            result_schema=[item.model_dump(mode="json") for item in columns],
@@ -702,7 +681,162 @@ def preview_pipeline(
        source_fingerprints=source_fingerprints,
        input_row_count=input_row_count,
        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,
    )