apache / apache/datafusion-python

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

Open
#1,719 0 comments 0 reactions 1 assignee Claimed by @timsaucer View on GitHub
Dominant language
Python
Stars
604
Forks
174
Avg merge
1d 7h
Merged PRs (30d)
4

Description

`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::()`, 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`. `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

- [ ] #1720 — Report physical partitioning, and stop two panics escaping as panics: G3, G4, G5, G8, and the comment for G2
- [ ] #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.

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.