A worked example of what datafusion-python's extension protocol is for: several independently compiled libraries, none of which knows about the others, cooperating on a single query whose work runs in separate operating system processes.
Everything here is real. The workers are separate interpreters. The plan they run was serialized by the driver and decoded by them. If you break the serialization, the tests fail.
| Crate | Owns | Installed with |
|---|---|---|
udf-library (dfx_udfs) | a scalar function, an aggregate, a window function | by hand — register_udf plus two with_*_extension_codec calls |
storage-library (dfx_storage) | a Parquet table provider and its own scan node | with_extensions |
engine-library (dfx_engine) | a query planner, a stage node, and the driver/worker machinery | with_extensions |
One of them is deliberately old-fashioned. dfx_udfs exposes no __datafusion_session_components__, so it cannot be installed as a bundle and its caller has to do five things in the right order instead of one. That is not a strawman: SessionExtensionComponents carries codec fields only, so a library that contributes functions has nowhere to put them today. Mixed setups are the normal case, and this example shows what one costs.
$ cd examples/distributed/engine-library $ uv venv $ uv pip install pytest maturin ../../.. ../storage-library ../udf-library $ uv run maturin develop $ uv run pytest
../../.. is the repository root. The example needs the datafusion built from this checkout rather than a released wheel, because an extension library and its host have to agree on the FFI ABI.
The other two libraries are installed for a reason beyond running their own tests: dfx_engine.session imports both, so import dfx_engine fails without them. That first uv pip install builds three Rust libraries from source and is not quick; maturin develop then builds the fourth, this one, in place.
Each library's tests can be run from its own directory the same way, without the sibling installs. storage-library and udf-library need only pytest maturin ../../...
Against real TPC-H data, which the script generates for itself:
$ uv pip install tpchgen-cli $ uv run python ../run_tpch.py --partitions 4
tpchgen-cli shards natively, so --partitions 4 asks it for four Parquet files and tells the engine to use four partitions — the two are the same number because a distributed engine can only spread work as widely as the data is split. --scale sets how much data; it defaults to a tenth of TPC-H scale factor 1, and --scale 1 is the full ~6M row lineitem.
The engine's planner splits the plan at the partial aggregate, which is where DataFusion has already split it for its own reasons: a GROUP BY becomes a partial pass per input partition and a final pass that merges them, and the partial passes are independent by construction.
SortPreservingMergeExec
ProjectionExec
AggregateExec: mode=FinalPartitioned <- driver merges
RepartitionExec: Hash([l_returnflag], 2)
FFI_ExecutionPlan: ShuffleStageExec <- shipped to workers
AggregateExec: mode=Partial <- one worker per partition
FFI_ExecutionPlan: PartitionedParquetExec
The driver serializes the ShuffleStageExec subtree, starts one worker per partition, and waits. Each worker rebuilds an equivalent session, decodes the plan, runs its partition, and writes the result to an Arrow IPC file. The driver then runs the whole query itself — and the stage node, finding the files already there, streams them instead of recomputing.
One node does both halves of that exchange, which is why nothing has to rewrite the plan in between. It also means a query run with no workers at all still gets the right answer; it just does the work itself.
The cost of that trick is that the filesystem is the state, and stage ids restart at 1 for every plan — so a second query pointed at the same shuffle directory would find the first one's files and read them. The driver refuses a directory that already holds stage output rather than letting that through: one shuffle directory per query.
engine-library/python/dfx_engine/session.py is the point of the whole example. There is no way to snapshot a SessionContext and restore it elsewhere, so worker parity cannot be automated — it has to be built the same way twice, from data small enough to put in a message. Both the driver and every worker call one build_session. Anything a query depends on that is not in the SessionSpec is a bug waiting for a worker to find it.
storage-library/src/codec.rs is the repository's only codec that encodes durable metadata. The others park the live object in a process-global map and encode an integer token, which is fine for making Rust type identity observable in a test and useless the moment the bytes leave the process. This one writes the file paths, the projection, and the schema, so the same bytes decode twice, decode on ten workers, and decode tomorrow.
udf-library/python/tests/_test_udfs.py shows that installing a library's codec is an alternative to registering its functions, not an addition. Three workers, three configurations:
| worker has | result |
|---|---|
| the codec, no registrations | works; the codec rebuilds each function from its name |
| the registrations, no codec | works; the registry answers first and the codec is never consulted |
| neither | fails, naming dfx_net_revenue |
The middle row is the trap. On the driver, where the functions are always registered, a broken or missing codec looks completely fine.
engine-library/python/tests/_test_three_libraries.py runs the queries, and pins the failure modes next to the successes — including a Python UDF that works on the driver and fails on the worker.
It writes shuffle results to local files, so “distributed” means several processes on one machine. Adding a network is a transport change and would not alter anything above it.
It holds one partition of results in memory before writing, because an Arrow IPC stream needs its schema up front. A production engine would stream to the file and track completion separately.
It has one level of stages. A query with an aggregate in each branch of a union gets one stage per branch, and they run side by side — but no stage ever feeds another. A real engine chains them, and the interesting problems — scheduling, retries, straggler handling, memory limits — all live in the part this example replaces with subprocess.Popen and a for loop.
It is slower than running the query in one process. Four processes on one laptop cannot beat one process that skips a round trip through Arrow IPC files. The comparison the tests make is agreement, not speed.