Add dfx_engine: a toy distributed engine in both halves

Third library for #1719, and the one that makes the other two do something.
Rust owns the query planner, the stage node, its codec, and a config
extension; Python owns the session factory, the driver, and the worker entry
point. A real engine needs both, so this crate is a mixed maturin package
rather than a pure extension module.

The split point is the partial aggregate. DataFusion already breaks a GROUP BY
into a partial pass per input partition and a final pass that merges them, so
the partial passes are independent by construction and only their output has
to come back. Wrapping that subtree in a `ShuffleStageExec` is the whole
rewrite.

The planner plans against `LocalOptimizerSession`, which borrows the foreign
session but owns the stock optimizer rule list. Without it the rules run back
across FFI and hand the library `ForeignExecutionPlan` wrappers, which cannot
be serialized (that is G1) and cannot be rewritten either — an engine cannot
split a subtree it holds only an opaque handle to. This was validated as a
spike before any of it was built.

One node does both halves of the shuffle. `execute(i)` reads the file for
partition `i` if it exists and otherwise computes its child and writes it on
the way past, so the same node is the thing a worker runs and the thing the
driver reads, and nothing has to rewrite the plan in between. A query with no
workers still gets the right answer, having done the work itself. The shuffle
directory travels inside the node and therefore inside its encoding, so a
worker and a driver cannot disagree about where results go.

`session.py` is the piece the whole example exists to motivate. There is no
way to snapshot a SessionContext and restore it elsewhere — SessionConfig is
write-only from Python, and `df_settings` is readable but lists
`datafusion.runtime.*` keys that have no namespace to set them back into — so
worker parity cannot be automated. It has to be built the same way twice from
data small enough to put in a message, which is what `SessionSpec` is. Both
sides call `build_session`; anything a query depends on that is not in the
spec is a bug waiting for a worker to find it.

Two findings this turned up, both now documented in the code:

`dfx_storage` needed a *logical* codec, not just a physical one. Its scan node
is physical, so a physical codec looks sufficient — but installing any FFI
query planner means the session hands that planner the logical plan as
protobuf, and a logical plan holds its tables as `Arc<dyn TableProvider>`. With
no `try_encode_table_provider` the session fails at `execution_plan()` with
"Error serializing custom table", before anything is distributed. The payload
is the directory, since everything else the provider holds is read back from
it.

A foreign node does not print its own name. The host shows
`FFI_ExecutionPlan: ShuffleStageExec, number_of_children=1`, so the driver's
tree walk has to match on containment; an anchored match works in a
single-library test and fails the moment a real extension is involved.

Verified end to end: three Parquet files, three worker processes, each
producing one partition of partial aggregate, driver merging them to the same
answer the single-process path gives. Corrupting one shuffle file breaks the
driver's query, which is how we know the workers did the work rather than the
driver quietly recomputing it.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: timsaucer

Cargo.lock

@@ -34,6 +34,7 @@ members = [
   "examples/datafusion-ffi-query-planner-example",
   "examples/distributed/storage-library",
   "examples/distributed/udf-library",
+  "examples/distributed/engine-library",
 ]
 resolver = "3"
 

Cargo.toml

@@ -0,0 +1,51 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+[package]
+name = "dfx-engine"
+version.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+license.workspace = true
+description = "Example extension library: a toy distributed engine that splits plans into stages"
+homepage.workspace = true
+repository.workspace = true
+publish = false
+
+[dependencies]
+arrow = { workspace = true }
+async-trait = { workspace = true }
+datafusion = { workspace = true }
+datafusion-common = { workspace = true, default-features = false }
+datafusion-ffi = { workspace = true }
+datafusion-proto = { workspace = true }
+datafusion-python-util.workspace = true
+datafusion-session = { workspace = true }
+futures = { workspace = true }
+pyo3 = { workspace = true, features = [
+  "extension-module",
+  "abi3",
+  "abi3-py310",
+] }
+pyo3-log = { workspace = true }
+
+[build-dependencies]
+pyo3-build-config = { workspace = true }
+
+[lib]
+name = "_internal"
+crate-type = ["cdylib", "rlib"]

examples/distributed/engine-library/Cargo.toml

@@ -0,0 +1,20 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+fn main() {
+    pyo3_build_config::add_extension_module_link_args();
+}

examples/distributed/engine-library/build.rs

@@ -0,0 +1,34 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+[build-system]
+requires = ["maturin>=1.6,<2.0"]
+build-backend = "maturin"
+
+[project]
+name = "dfx_engine"
+requires-python = ">=3.10"
+classifiers = [
+  "Programming Language :: Rust",
+  "Programming Language :: Python :: Implementation :: CPython",
+]
+dynamic = ["version"]
+
+[tool.maturin]
+features = ["pyo3/extension-module"]
+python-source = "python"
+module-name = "dfx_engine._internal"

examples/distributed/engine-library/pyproject.toml

@@ -0,0 +1,40 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""A toy distributed engine, as an extension library.
+
+Two halves, because a real engine has two: the Rust side owns the query
+planner, the stage node, the codec that carries it, and a config extension;
+the Python side owns the session factory, the driver, and the worker entry
+point.
+
+Start with :mod:`dfx_engine.session` -- ``build_session`` is the piece the
+rest of the example exists to motivate.
+"""
+
+from dfx_engine import _internal
+from dfx_engine._internal import DfxEngineConfig, DfxEngineExtension
+from dfx_engine.session import SessionSpec, build_session, expected_codec_ids
+
+__all__ = [
+    "DfxEngineConfig",
+    "DfxEngineExtension",
+    "SessionSpec",
+    "_internal",
+    "build_session",
+    "expected_codec_ids",
+]

examples/distributed/engine-library/python/dfx_engine/__init__.py

@@ -0,0 +1,206 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""The driver: split a query into tasks, fan them out, collect the answer.
+
+The shape is deliberately boring, because the interesting part is not the
+scheduling. What matters is the four things the driver has to get right, each
+of which is a way a real deployment goes wrong:
+
+1. It serializes the stage **with** its session. ``to_bytes(None)`` uses an
+   empty codec chain and cannot encode any library's node.
+2. It ships the codec ids it used, so a worker can refuse a plan it would
+   misread rather than decode it with the wrong codec.
+3. It puts the shuffle directory in the session config, not in the message,
+   so the directory travels *inside* the encoded plan and the two sides
+   cannot disagree.
+4. It waits for every worker before reading, because the stage node decides
+   whether to read or recompute by looking at the filesystem.
+"""
+
+from __future__ import annotations
+
+import json
+import pathlib
+import subprocess
+import sys
+from typing import TYPE_CHECKING
+
+from dfx_engine import _internal
+from dfx_engine.session import SessionSpec, build_session
+
+if TYPE_CHECKING:
+    import pyarrow as pa
+    from datafusion import DataFrame, SessionContext
+    from datafusion.plan import ExecutionPlan
+
+__all__ = ["DistributedResult", "find_stage", "run_distributed"]
+
+
+class DistributedResult:
+    """What a distributed run produced, and how."""
+
+    def __init__(
+        self,
+        batches: list[pa.RecordBatch],
+        partitions: list[int],
+        worker_rows: dict[int, int],
+    ) -> None:
+        self.batches = batches
+        self.partitions = partitions
+        """Partition indices that were dispatched, one per worker."""
+        self.worker_rows = worker_rows
+        """Rows each worker produced, keyed by partition index."""
+
+
+STAGE_NODE_NAME = "ShuffleStageExec"
+
+
+def find_stage(plan: ExecutionPlan) -> ExecutionPlan | None:
+    """Locate the stage node the planner inserted.
+
+    Matched on the display string because a Python caller has no way to
+    downcast a Rust plan node -- there is no ``isinstance`` across an FFI
+    boundary.
+
+    Note the *containment* test. The node was built inside this library and
+    handed back to the host, so what the host prints is not
+    ``ShuffleStageExec: stage=1`` but::
+
+        FFI_ExecutionPlan: ShuffleStageExec, number_of_children=1
+
+    A foreign node reports its own name nested inside the wrapper's, which
+    makes anchored matches on plan text quietly wrong -- the kind of thing
+    that works in a single-library test and fails the moment a real extension
+    is involved.
+    """
+    if STAGE_NODE_NAME in plan.display():
+        return plan
+    for child in plan.children():
+        found = find_stage(child)
+        if found is not None:
+            return found
+    return None
+
+
+def _dispatch(
+    envelope: dict, envelope_dir: pathlib.Path, partition: int
+) -> subprocess.Popen[str]:
+    """Start one worker for one partition.
+
+    ``sys.executable``, not ``python``: a worker on a different Python minor
+    version cannot load a cloudpickled inline UDF, and that failure is far
+    from its cause.
+    """
+    path = envelope_dir / f"task-{partition}.json"
+    path.write_text(json.dumps(envelope))
+    return subprocess.Popen(  # noqa: S603
+        [sys.executable, "-m", "dfx_engine.worker", str(path)],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+    )
+
+
+def run_distributed(sql: str, spec: SessionSpec) -> DistributedResult:
+    """Run `sql`, executing its leaf stage in one worker process per partition.
+
+    Requires ``spec.shuffle_dir``: without it the planner inserts no stage and
+    there is nothing to distribute.
+    """
+    if not spec.shuffle_dir:
+        message = "run_distributed needs a shuffle_dir; build_session got none"
+        raise ValueError(message)
+
+    ctx, engine, _storage = build_session(spec)
+    plan = ctx.sql(sql).execution_plan()
+
+    stage = find_stage(plan)
+    if stage is None:
+        message = (
+            "no ShuffleStageExec in the plan; the engine's planner did not run, "
+            "or its config extension was not registered"
+        )
+        raise RuntimeError(message)
+
+    shuffle_dir = pathlib.Path(spec.shuffle_dir)
+    shuffle_dir.mkdir(parents=True, exist_ok=True)
+
+    # Encode the stage subtree, through the session that owns the codecs.
+    plan_path = shuffle_dir / "stage.plan"
+    plan_path.write_bytes(stage.to_bytes(ctx))
+
+    stage_id = _internal.stage_id()
+    partitions = list(range(stage.partition_count))
+    envelopes = [
+        {
+            "spec": spec.to_json(),
+            "plan": str(plan_path),
+            "stage_id": stage_id,
+            "partition": partition,
+        }
+        for partition in partitions
+    ]
+
+    # One process per partition, all in flight together. This is the claim the
+    # example is making: each worker reads a different file and writes a
+    # different result, so they need no coordination beyond the directory.
+    workers = [
+        _dispatch(envelope, shuffle_dir, partition)
+        for envelope, partition in zip(envelopes, partitions, strict=True)
+    ]
+
+    worker_rows: dict[int, int] = {}
+    failures = []
+    for partition, worker in zip(partitions, workers, strict=True):
+        stdout, stderr = worker.communicate()
+        if worker.returncode != 0:
+            failures.append(f"partition {partition} failed:\n{stderr}")
+            continue
+        worker_rows[partition] = json.loads(stdout)["rows"]
+
+    if failures:
+        raise RuntimeError("\n".join(failures))
+
+    # Now run the whole query here. Every stage partition has a file, so the
+    # stage node streams them instead of recomputing -- the driver does only
+    # the final merge.
+    batches = ctx.sql(sql).collect()
+    _ = engine
+    return DistributedResult(batches, partitions, worker_rows)
+
+
+def run_local(sql: str, spec: SessionSpec) -> list[pa.RecordBatch]:
+    """Run `sql` in this process, for comparison.
+
+    Uses the same session factory with no shuffle directory, so the only
+    difference from :func:`run_distributed` is where the work happened.
+    """
+    ctx, _engine, _storage = build_session(
+        SessionSpec(
+            tables=spec.tables,
+            shuffle_dir="",
+            target_partitions=spec.target_partitions,
+        )
+    )
+    return ctx.sql(sql).collect()
+
+
+def dataframe_for(sql: str, spec: SessionSpec) -> tuple[SessionContext, DataFrame]:
+    """Session and DataFrame for `sql`, for tests that want to inspect a plan."""
+    ctx, _engine, _storage = build_session(spec)
+    return ctx, ctx.sql(sql)