Add FFI query planner support by timsaucer · Pull Request #1677 · apache/datafusion-python

timsaucer · 2026-08-07T19:21:15Z

Which issue does this PR close?

Related to #1612. This PR does not close it, but provides the FFI query planner plumbing that a datafusion-distributed integration can build on.

This is part 1 of 3 in the split of #1672. These are enabled as a github stack so you should be able to swab between the 3 PRs in github interface (above, next to the "Open" oval).

Rationale for this change

Extension libraries (for example distributed execution engines) need to supply their own QueryPlanner to a SessionContext without compiling against the datafusion-python crate. This PR exposes the query planner over the FFI boundary, following the same PyCapsule pattern used for table providers and catalogs.

What changes are included in this PR?

SessionContext.with_query_planner(planner) installs a planner exported via a __datafusion_query_planner__ PyCapsule, preserving existing session state and codec settings.
SessionContext.__datafusion_query_planner__() exports the current planner so another planner can wrap it as an explicit fallback (a session holds exactly one planner; layering is explicit delegation).
The Python codecs keep the exporting SessionContext alive for the capsule getters, since FFI_TaskContextProvider holds its provider weakly and rebinding ctx would otherwise kill a capsule captured from it.
New example crate datafusion-ffi-query-planner-example demonstrating a real three-library plan exchange (host, provider library, planner library as separate cdylibs), including session config transfer via SessionConfig.with_extension.
New docs/source/contributor-guide/ffi.md sections covering the query planner capsule protocol, what a derived context shares, the task context provider a planner exports, and the three-library setup.

Are there any user-facing changes?

New public APIs: SessionContext.with_query_planner and SessionContext.__datafusion_query_planner__. A new example crate ships under examples/. No breaking changes to existing APIs.

AI Disclosure: This code was written in part by an AI agent.:

timsaucer · 2026-08-07T19:35:04Z

        Ok(())
    }

+    pub fn with_query_planner(&self, planner: Bound<'_, PyAny>) -> PyDataFusionResult<Self> {


This API is the main reason for this PR. Here we allow changing out the default query planner with a user provided query planner.

timsaucer · 2026-08-07T19:36:39Z

+      - name: Build FFI query planner test library
+        if: matrix.python-tag == 'abi3'
+        uses: PyO3/maturin-action@v1
+        with:
+          target: x86_64-unknown-linux-gnu
+          manylinux: "2_28"
+          working-directory: examples/datafusion-ffi-query-planner-example
+          args: --out dist
+          rustup-components: rust-std


In order to prove that the 3 library approach works where we have different codecs and different execution plans provided, we are adding a second test library. This way we can make sure there is no accidental ability to reach into a foreign code block.

timsaucer · 2026-08-07T19:37:19Z

+struct RuntimeAwareQueryPlanner {
+    planner: FFI_QueryPlanner,
+}


As the docstring says, the purpose of this is to make sure we attach the runtime handle when needed.

timsaucer · 2026-08-07T19:38:38Z

+    pub fn __datafusion_query_planner__<'py>(
+        &self,
+        py: Python<'py>,
+    ) -> PyResult<Bound<'py, PyCapsule>> {


We need our session context to export it's own query planner because we have a use case where one query planner can depend on another. This is already supported by datafusion-distributed, so we want to be certain we support it here.

timsaucer · 2026-08-07T19:41:57Z

+#[derive(Clone, Debug)]
+pub(crate) struct PlannerConfig {
+    pub max_rows: usize,
+}


I'm adding this to the query planner example because it's a very common pattern that we will need custom configs for the query planner, so it is reasonable to need insurance that configs pass over the FFI boundary properly and to use as a demonstration to anyone who is providing such a library.

this is needed for ballista, thanks Tim for example

The FFI test wheel artifact now bundles two projects, so upload-artifact preserves a `<project>/dist/` prefix instead of placing the wheels at the artifact root. The install step globbed `wheels/*.whl`, which no longer matched them, so the FFI wheels were silently skipped and the FFI unit tests failed with `ModuleNotFoundError: No module named 'datafusion_ffi_example'`. Install the recursive `find` results instead of re-globbing. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

ntjohnson1

Appears consistent with the rest of the FFI plumbing

ntjohnson1 · 2026-08-14T18:13:07Z

        """
        self.ctx.add_physical_optimizer_rule(rule)

+    def with_query_planner(


Generally wonder if this builder pattern feels pythonic. Consistent with what's already here so no action requested. Didn't look at how many withs there are but

ctx = SessionContext(config, planner)

feels a little more intuitive than

ctx = SessionContext().with_query_planner(planner)

Good point! Also worth updating the skill to match this pattern

milenkovicm

thanks @timsaucer cant want to get this integrated

milenkovicm · 2026-08-17T11:13:59Z

+}
+
+#[pymethods]
+impl PlannerConfig {


Nit, MyPlannerConfig to have names aligned,

milenkovicm · 2026-08-17T11:14:18Z

+#[derive(Clone, Debug)]
+pub(crate) struct PlannerConfig {
+    pub max_rows: usize,
+}


this is needed for ballista, thanks Tim for example

milenkovicm · 2026-08-17T11:20:32Z

+            observations: Arc::clone(&self.observations),
+        });
+        let runtime = get_tokio_runtime().handle().clone();
+        let ctx_provider = Arc::new(SessionContext::new()) as Arc<dyn TaskContextProvider>;


is this session context be parameter of method call on the line 119 ? are those two different sessions ?

Collapse the two duplicated planner-install blocks into a single `ctx_with_rebound_planner`. A derived context shares the existing `SessionContext` when there is no foreign planner to rebind, and forks only when one is installed, since the FFI codecs capture the context they are built against. Document what that fork shares. Catalogs, tables, and the runtime environment stay shared; registered functions, configuration, and the optimizer rule lists are snapshotted. The caveat lands on all four derivation methods and on a new contributor-guide subsection, with tests covering both halves. Explain why `RuntimeAwareQueryPlanner` exists at all. Upstream's `ForeignQueryPlanner` is the consumer-side adapter that lets an `FFI_QueryPlanner` satisfy the `QueryPlanner` trait, which is what makes a planner from another shared library installable in a `SessionState`. Its trait method receives only a `&LogicalPlan` and a `&dyn Session`, so it has nowhere to obtain a runtime handle and passes `None`. Throughout datafusion-ffi each library attaches its own runtime to the objects it exports, so a producer-side wrapper can enter that runtime before running its own library's code. A provider owned by another library keeps its owner's runtime even when it travels through our catalog, because `FFI_TableProvider::new_with_ffi_codec` unwraps a `ForeignTableProvider` back to the original handle and discards the runtime passed alongside it. `session_runtime` is that same rule applied to the session: `FFI_SessionRef` is our object and every callback on it runs our code. It matters for what those callbacks hand back. A plan produced by our own planner returns as `FFI_ExecutionPlan::new(plan, runtime)`, and `execute` enters that runtime before calling into the plan; the same holds for our physical optimizer rules and for tables we own rather than re-export. The delegation case this type exists for is exactly that shape. A foreign planner falling back to our planner through `__datafusion_query_planner__` receives a plan whose execution needs our runtime, and datafusion-python owns that runtime as a process global while the Python thread calling in carries no ambient one. The same reasoning is why `__datafusion_query_planner__` re-exports through the adapter rather than unwrapping to the inner handle. A consumer reaching us through `ForeignQueryPlanner` calls with `None`, so the adapter is what restores our handle on the way back out. Unwrapping would save a planning-time round trip and silently drop it. In the planner example, match the two real spellings of the row-limit config key exactly instead of by suffix, and validate after both lookup paths so the fallback cannot accept `max_rows = 0`. The key appears twice because rebuilding a `ConfigOptions` across the FFI boundary parks every foreign extension inside a single `FFI_ExtensionOptions`, itself namespaced under `datafusion_ffi`. Also declare `requires-python = ">=3.10"` on the provider example to match the `abi3-py310` feature it builds against, and link both example READMEs to the contributor guide rather than restating its caveats. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Remove `RuntimeAwareQueryPlanner`. It existed to re-attach our Tokio handle to the session we hand to a foreign planner, on the reasoning that `ForeignQueryPlanner` passes `session_runtime: None`. That handle turns out to have no reachable path: the query planner FFI exchanges serialized bytes rather than plan handles, a provider owned by another library keeps its own runtime because `FFI_TableProvider::new_with_ffi_codec` unwraps a `ForeignTableProvider` back to the original handle, and we execute on our own runtime regardless. Setting the handle to `None` left every test passing. Codec rebinding now downcasts upstream's `ForeignQueryPlanner` directly, which also stops `__datafusion_query_planner__` adding a second layer, since `new_with_ffi_codecs` already unwraps that type. The `datafusion-session` dependency is no longer needed in crates/core. Keep the exporting session alive for codecs handed out in a PyCapsule. `FFI_TaskContextProvider` stores its provider in a `Weak`, so a capsule stopped working as soon as the `SessionContext` that produced it went out of scope. That made the natural spelling of the documented fallback pattern fail: fallback = ctx.__datafusion_query_planner__() ctx = ctx.with_query_planner(MyPlanner(fallback=fallback)) Rebinding `ctx` dropped the exporter and planning then failed with "TaskContextProvider went out of scope over FFI boundary". Both Python codecs gained an opt-in `exported_session`, set only by the three capsule getters. The keep-alive lives in the inner codec because the consumer clones the FFI handle out of the capsule and `clone` clones the inner codec's `Arc`, so a capsule-scoped keep-alive would die too early. It is deliberately opt-in: the same codecs are also attached to providers and catalogs that end up back inside the session, where a strong reference would close a `SessionContext -> SessionState -> query planner -> FFI codec` cycle. Both structs now implement `Debug` by hand, because `SessionContext` is not `Debug`. Add two example tests. One drives a plan containing `RepartitionExec`, which spawns Tokio tasks as it runs, through all three libraries, so the codecs are exercised on a multi-node plan rather than a bare scan. The other layers a planner on top of the session's existing planner using the capsule captured beforehand, which is the delegation pattern upstream prescribes; `Session::create_physical_plan` cannot be used for this, because it dispatches through the installed planner and recurses. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…m fix

timsaucer added 4 commits August 6, 2026 10:54

Add FFI query planner support

9ba147c

AI Disclosure: This code was written in part by an AI agent.:

Add three-library FFI planner example

47aea03

AI Disclosure: This code was written in part by an AI agent.:

Update FFI query planner integration

05ab5da

AI Disclosure: This code was written in part by an AI agent.:

add rat

27147f3

timsaucer mentioned this pull request Aug 7, 2026

Add FFI query planner support with composable extension codecs #1672

Closed

timsaucer commented Aug 7, 2026

View reviewed changes

ntjohnson1 approved these changes Aug 14, 2026

View reviewed changes

Merge branch 'main' into feat/ffi-query-planner-core

c744579

milenkovicm approved these changes Aug 17, 2026

View reviewed changes

timsaucer mentioned this pull request Aug 26, 2026

Session::create_physical_plan over FFI ignores the session's LogicalExtensionCodec apache/datafusion#24688

Closed

timsaucer and others added 4 commits August 26, 2026 10:48

add override for datafusion version to pre-release testing of upstrea…

c6ba98f

…m fix

remove unintentionally committed files

752243b

Empty commit to trigger CI

e9dc22c

timsaucer mentioned this pull request Aug 26, 2026

Release datafusion-python 55.0.0 #1696

Open

Navigation Menu

Uh oh!

Add FFI query planner support - #1677

Add FFI query planner support#1677
timsaucer wants to merge 11 commits into
mainfrom
feat/ffi-query-planner-core

timsaucer commented Aug 7, 2026 •

edited

Loading

Uh oh!

timsaucer Aug 7, 2026

Uh oh!

timsaucer Aug 7, 2026

Uh oh!

timsaucer Aug 7, 2026

Uh oh!

timsaucer Aug 7, 2026

Uh oh!

timsaucer Aug 7, 2026

Uh oh!

milenkovicm Aug 17, 2026

Uh oh!

ntjohnson1 left a comment

Uh oh!

ntjohnson1 Aug 14, 2026

Uh oh!

timsaucer Aug 15, 2026

Uh oh!

milenkovicm left a comment

Uh oh!

milenkovicm Aug 17, 2026

Uh oh!

milenkovicm Aug 17, 2026

Uh oh!

milenkovicm Aug 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

timsaucer commented Aug 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are there any user-facing changes?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ntjohnson1 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

milenkovicm left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

timsaucer commented Aug 7, 2026 •

edited

Loading