Add a first-class example of multiple extension libraries in one distributed query

[timsaucer]

with_extensions (#1679) and composable codecs (#1678) gave extension libraries a way to ship codecs and a query planner as one atomic bundle. What the repository still has no example of is the thing that machinery exists for: several independently-compiled libraries cooperating inside one distributed query, with plans actually leaving the process.

This issue tracks building that example, and records the gaps found while validating that it is possible. Several of them are load-bearing, and one is an upstream blocker that the current examples only avoid by accident.

The example

Three extension libraries under examples/distributed/, plus a toy distributed engine that runs real OS-process workers:

• A UDF library — scalar UDF, UDAF and UDWF resolved by name, with a logical codec whose payload is empty. Deliberately does not implement __datafusion_session_components__, so it is installed the old way with register_udf plus with_logical_extension_codec. This is the mixed-workflow case, and it is currently the honest one: SessionExtensionComponents has no UDF field yet.
• A storage library — a Parquet-directory table provider reporting one output partition per file, with a physical codec that emits genuinely portable bytes rather than a process-local token. There is no such codec anywhere in the repository today (see G9).
• An engine library — stage-splitting query planner, its own StageExec and ShuffleReadExec nodes, both codecs, and a mixed Rust/Python maturin package carrying the driver, the worker entry point, and the shared session factory.

Workers are separate processes spawned through sys.executable; results come back as Arrow IPC files, one per (stage, partition). Queries run against TPC-H lineitem at SF 1, which CI already generates. One query uses an inline Python UDF defined in the driver's __main__, paired with the case that breaks — the same UDF defined in an importable module the worker cannot import.

The existing datafusion-ffi-query-planner-example is retired into the engine library; datafusion-ffi-example stays as the protocol-conformance test bed.

Gaps found

Every item below was reproduced against a built extension, not inferred.

One correction to how this was originally framed: the datafusion-python items were found while building the example, not required by it. The finished example depends on none of them — it reads partition_count, never calls SessionConfig.set, and bounds-checks partitions in Python before reaching the Rust path. #1720 stands on the defects themselves.

Blocked upstream

G1 — a stock node that crosses FFI cannot be serialized on the query-planner return path. FFI_QueryPlanner returns proto bytes rather than a plan handle, so every query serializes the physical plan. The plan contains a CooperativeExec inserted by the always-on EnsureCooperative rule, which runs on the host during a foreign planner's create_physical_plan and so arrives inside the library as an opaque ForeignExecutionPlan. That type has no reachable try_to_proto, so the native encoder never runs and a perfectly serializable node becomes unserializable purely by having crossed the boundary.

G2 — the greedy codec claim in datafusion-ffi-example is a symptom of G1, not sloppiness. Its physical codec claims node.is::<ForeignExecutionPlan>(), which takes every other library's nodes, and extension-guide/checklist.md tells authors never to do this. It is nevertheless load-bearing: narrowing it to DataSourceExec alone makes 31 of the 51 tests in datafusion-ffi-query-planner-example fail, every one on the CooperativeExec node from G1. The arm has to stay until G1 is fixed. A planner that controls its own physical optimizer rules never sees a foreign node and needs no such arm — which is how the new engine library will avoid it.

G7 — a codec reached over FFI gets a TaskContext with no object stores and no catalog. FFI_TaskContext is rebuilt with RuntimeEnv::default(), so decode-time object-store resolution — which ParquetSource::try_from_proto performs — can only ever see file://, and no codec can resolve a table by name at decode time.

datafusion-python

G3 — no way to read a plan's partitioning scheme. Only partition_count was exposed, so nothing in Python could distinguish hash-distributed output from merely counted output, or read the hash keys. Fixed by ExecutionPlan.output_partitioning.

Worth being accurate about the motivation, since I first recorded this as blocking: the example does not need it. Its engine decides the split in Rust, inside the planner, where output_partitioning() was always reachable — so the Python accessor was never on the critical path. What it does buy is the ability to assert that a scan and a stage report UnknownPartitioning rather than Hash, which matters because claiming a hash partitioning a plan does not have would let the optimizer skip a repartition it needs.

G4 — SessionContext.execute did not bounds-check the partition index. The plan's leaves index their partition vector directly, so an out-of-range index panicked and surfaced as index out of bounds: the len is 2 but the index is 5, naming neither the plan nor the index. Fixed.

G5 — SessionConfig.set raised PanicException. It routed through SessionConfig::set_str, which unwraps, so an unknown namespace aborted rather than raised — and PanicException derives from BaseException, escaping except Exception. Fixed on the Python side; the upstream unwrap remains.

G6 — there is no session-config snapshot or restore, and nothing documents what a worker has to match. SessionConfig and RuntimeEnvBuilder are write-only from Python. information_schema.df_settings is readable but not replayable: it lists datafusion.runtime.* keys that have no ConfigOptions namespace, so a naive replay loop hits G5 on its first row. Worker parity has to be hand-maintained, and the checklist for doing so does not exist yet.

G11 — a Python worker cannot learn a decoded plan's schema. Neither ExecutionPlan.schema() nor RecordBatchStream.schema is exposed, so Python cannot write a correctly-typed empty result for a plan it just decoded. Worked around in the example by having the Rust stage node write its own shuffle file, which is better design regardless, but the accessor is a real omission.

G12 — installing any FFI query planner forces every table provider in the session to be logically encodable. The planner receives the logical plan as protobuf, and a logical plan holds its tables as Arc<dyn TableProvider>. DefaultLogicalExtensionCodec::try_encode_table_provider is not_impl_err!, so a provider library shipping only a physical codec — which looks sufficient, since its scan node is a physical node — fails at execution_plan() with Error serializing custom table the moment an engine is installed, before anything is distributed. Found by building it that way; the storage library now ships both codecs. This is protocol behaviour rather than a bug, and it is not written down anywhere.

G13 — a foreign node does not display its own name. The host prints FFI_ExecutionPlan: ShuffleStageExec, number_of_children=1, so a Python-side plan-text match has to test containment rather than prefix. An anchored match passes in a single-library test and fails as soon as a real extension is involved.

G10 — the example planner emits an invalid plan on multi-partition input. DistributedQueryPlanner wraps GlobalLimitExec after optimization, so nothing inserts the coalesce it requires: Assertion failed: self.input.output_partitioning().partition_count() == 1 (left: 2, right: 1). Latent because every test for it uses a single-partition input.

Documentation

G8 — two ExecutionPlan docstrings claimed memory-backed tables cannot be serialized. True of LogicalPlan, whose try_encode_table_provider has no arm for one; false of the physical layer, which inlines the batches. Verified by decoding on a context sharing nothing with the encoder and executing. Fixed.

G9 — extension_codec_durable_metadata has no reference implementation. The guide tells authors to encode durable metadata and states that the in-repo examples deliberately do not. Nothing in the repository shows one that does. The new storage library becomes that reference.

Validated, for the record

The design depends on these, so each was checked rather than assumed: Parquet and memory-backed physical plans round-trip through to_bytes/from_bytes on a fresh context with no codecs installed; partition_count survives the FFI boundary; a __main__-defined Python UDF ships by value to a genuinely unrelated process; codec-id dispatch is order-independent and names the missing id on failure; and a query planner that wraps the foreign session with its own physical optimizer rule list produces a plan free of ForeignExecutionPlan, which another process then decodes and executes.

PRs

• Report physical partitioning and resolve two panics escaping as PanicException #1720 — Report physical partitioning, and stop two panics escaping as panics: G3, G4, G5, G8, and the comment for G2
• Add a worked example of three extension libraries in one distributed query #1721 — The multi-library distributed example itself: G6, G9, G12, G13, and the worker-parity documentation
• Upstream issues against apache/datafusion for G1, G5's unwrap, and G7

Still open here

G3, G4, G5 (Python side) and G8 are fixed by #1720. G6, G9, G12 and G13 are addressed by documentation and by the example in #1721. G10 is a latent bug in datafusion-ffi-query-planner-example, not yet fixed. G11 (ExecutionPlan.schema() / RecordBatchStream.schema) is unfixed and was worked around. G1, G2 and G7 need upstream changes.