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
@@ -18,6 +18,7 @@ from govoplan_core.core.datasources import (
    DatasourcePublicationResult,
 )
 from govoplan_core.db.base import Base
+from govoplan_dataflow.backend.backends.duckdb import DuckDbExecutionBackend
 from govoplan_dataflow.backend.db.models import (
    DataflowPipeline,
    DataflowPipelineRevision,
@@ -320,6 +321,37 @@ class DataflowServiceTests(unittest.TestCase):
            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:
        graph = sample_graph()
        graph.nodes[0].config["rows"] = [{"id": 1, "amount": "not-a-number"}]
@@ -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()