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 three capsule getters that need something only a session has — __datafusion_query_planner__, __datafusion_logical_extension_codec__, and __datafusion_physical_extension_codec__ — now receive the SessionContext they are being installed on, matching what __datafusion_table_provider__ and friends have done since 52.0.0. A codec takes the TaskContextProvider from it; a planner takes both codecs. Neither example constructs a SessionContext any more, and decode callbacks now resolve names against the session running the query.
PySessionContext retains forked ancestors, so a codec bound to a session before a planner fork cannot be left holding a dropped weak reference.
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.
Both example codecs accept require_udf_on_decode, and tests assert which session a decode callback resolves against — including a function registered on the host, one registered after the codec was installed, and the fork boundary.
New docs/source/contributor-guide/ffi.md sections covering the capsule protocol, what a derived context shares, and the fork caveat. New .ai/skills/ffi-capsule-protocol/ recording the convention, with a pointer from AGENTS.md. Corrected user-guide/io/table_provider.md, which still showed the pre-52.0.0 signature.

Are there any user-facing changes?

Yes, including a breaking change.

New public APIs: SessionContext.with_query_planner and SessionContext.__datafusion_query_planner__.

Breaking: __datafusion_logical_extension_codec__ and __datafusion_physical_extension_codec__ now take a session: Bound<PyAny> parameter, so any extension library implementing them must be updated. docs/source/user-guide/upgrade-guides.md has a 55.0.0 section with before and after. Calling the old signature raises an import error naming the method rather than a bare TypeError. SessionContext's own getters accept the argument optionally, so ctx.__datafusion_logical_extension_codec__() is unaffected.

A new example crate ships under examples/.

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 ?

Really good catch! This led me down a rabbit hole and I ended up needing two upstream fixes:

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

FFI constructors silently discard arguments when the input is already foreign datafusion#24722

In the latest push we no longer create this session context just for the codecs.

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

`FFI_TaskContextProvider` downgrades the provider it is given to a `Weak`, so building one inline in `__datafusion_query_planner__` left the capsule carrying a provider that was already dropped by the time it returned. Every codec callback through that capsule would have failed with "TaskContextProvider went out of scope over FFI boundary". The example did not notice because it ships the default codecs and no custom extension nodes, so `try_decode` is never reached. `MyQueryPlanner` now owns the context and hands out clones of it. The `QueryPlanner` the capsule carries holds a reference too, so the capsule stays usable even when the Python object that exported it is dropped first. Document the distinction the inline construction obscured. The `TaskContextProvider` supplied at export time backs the exporting library's own codec callbacks, decoding that library's nodes in its own registry. It is unrelated to the `&dyn Session` that later arrives at `create_physical_plan`, which belongs to the host, and it could not be derived from that session in any case, since the codecs are built before any session exists. Rename `PlannerConfig` to `MyPlannerConfig` to match `MyQueryPlanner`. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

The example codecs restore objects from a process-local token registry and never read the `TaskContext` their FFI decode callbacks are handed, so which session that context belongs to was untestable. The token path ignores the registry entirely, which is why an empty `SessionContext::new()` has served as the exported provider without anyone noticing. Both codecs now accept `require_udf_on_decode`. When set, every decode call resolves that scalar function out of the task context it was given and fails with the session id if it is absent, which makes the answer observable. Each codec registers a marker function on the context it exports, so a name owned by the codec's library and a name owned by the host can be told apart. Four tests use it. The two library-local cases pass: a foreign codec resolves against the session its own library supplied. The two host-registered cases are `xfail(strict=True)`, because a function registered on the host with `register_udf` is not visible to a foreign codec's decode callback at all. A fifth pins the current error so the failure mode stays legible. Strict xfail means the pair will announce itself if the upstream design changes. Document the rule this establishes, and correct the surrounding section: `with_query_planner` rebuilds a foreign planner against the session that will run the query, so the provider a planner library supplies is replaced on that path. Codecs installed through `with_logical_extension_codec` and `with_physical_extension_codec` keep the provider their own library exported, which is the case these tests exercise. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

`FFI_QueryPlanner::new` and `FFI_{Logical,Physical}ExtensionCodec::new` ask an extension library for a `TaskContextProvider`, and a planner for two codecs on top of that. A library has none of those. Both examples answered with `Arc::new(SessionContext::new())`, an empty session that resolves nothing, held weakly by `FFI_TaskContextProvider` and therefore also a lifetime hazard. The table provider protocol already solved this: the host calls `__datafusion_table_provider__(session)` and the library takes what it needs off the session. Do the same for the other three getters. `__datafusion_query_planner__`, `__datafusion_logical_extension_codec__`, and `__datafusion_physical_extension_codec__` now receive the `SessionContext` they are being installed on. A codec takes the task context provider from it; a planner takes both codecs and uses `new_with_ffi_codecs`, which needs no provider at all. Neither example constructs a `SessionContext` any more. Decode callbacks consequently resolve against the session running the query. The two `xfail(strict=True)` tests from the previous commit now pass unmodified: a scalar function registered on the host with `register_udf` is visible inside a decode callback executing in another library, for both the logical and physical codec. A negative control keeps the check honest, and a further test covers a function registered after the codec was installed, since the provider is a live handle rather than a snapshot. `PySessionContext` gains an `ancestors` list. A foreign codec is built against the session current at the time it is installed and holds it weakly, so installing a foreign planner afterwards — which forks — would strand the codec once the Python name is rebound. The keep-alive lives on `PySessionContext` rather than on the codec because nothing reachable from a `SessionContext` reaches a `PySessionContext`, so it cannot close a cycle. What it does not paper over is the fork itself: a function registered after the fork is not visible to a codec bound to the session before it, which is the existing derived-context caveat seen from the codec's side, and is covered by a test. `SessionContext` accepts and ignores the argument on all three getters, so a session satisfies the same protocol a library implements and `ctx.__datafusion_query_planner__()` keeps working for the delegation pattern. Calling a stale getter that takes no session now reports an incompatible-library error naming the method, matching what `table_provider_from_pycapsule` does. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

The session-passing rule was already settled for four getters and documented in the 52.0.0 upgrade guide, but nothing pointed an agent or a new contributor at it before they wrote a fifth. Write it down where it will be found. Add the 55.0.0 upgrade guide entry this branch owes. Changing `__datafusion_logical_extension_codec__` and `__datafusion_physical_extension_codec__` to take a session breaks every extension library implementing them, so it needs before/after Rust in the same shape as the 52.0.0 entry. Correct `user-guide/io/table_provider.md`. It still showed the pre-52.0.0 signature with no session and a `PyCapsule::new_bound` call, so the one page a reader is most likely to find contradicted the convention. Add `.ai/skills/ffi-capsule-protocol/`. Its description is written as a trigger rather than a task, because the existing skills are all things to run on request and a convention read as one would be skipped. It leads with enumerating the family, which is the step that makes the rest unnecessary. Point `CLAUDE.md` at it, since that file loads unconditionally and a skill only helps once someone goes looking. Also note that `docs/temp/` is gitignored build output that `grep -r` surfaces with stale copies, and require an upgrade guide section alongside the `api change` label, so a breaking change forces a visit to the file that records the conventions. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Every `.ai/skills/*/SKILL.md` opened with the ASF header and only then the YAML frontmatter, which has to be the first thing in the file. The result was that no skill's `description` was readable: the skill listing showed `<!---` for all of them, so the field meant to say when a skill applies said nothing. `skills/datafusion_python/SKILL.md` already had the right order and was the model to follow. Move the header below the frontmatter in all four. Apache RAT still approves each file — it looks for the license anywhere, not at the top — verified with rat 0.13. This matters most for the new `ffi-capsule-protocol` skill, whose description is written as a trigger condition rather than a task name. The existing skills are all tasks to run on request, so a convention that has to be read *before* writing code is easy to filter out while skimming for something to invoke. Note the distinction in the skills section of `AGENTS.md`. Then remove what that makes redundant. `AGENTS.md` had grown a copy of the skill's opening grep and a summary of its central rule. Two copies of one convention, with the more discoverable copy free to drift, is exactly the failure this branch already fixed in `user-guide/io/table_provider.md`. `AGENTS.md` now says only when to look and where; the skill owns the procedure. The `docs/source` versus `docs/temp` note moves the other way, out of the skill and into `AGENTS.md`, where it applies to everything rather than to this one protocol. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Installing a foreign query planner writes to `SessionState`, and `with_query_planner` must not modify its receiver, so it forks. A foreign codec holds an `FFI_TaskContextProvider` pointing at the session it was installed on, and until now the fork could not move it: passing a new provider to `FFI_LogicalExtensionCodec::new` was silently discarded whenever the codec was already foreign. The fork rebound only its own outer wrapper, so decode callbacks in the extension library kept answering from the pre-fork registry, and the pre-fork session had to be retained or the weakly held provider dangled. apache/datafusion#24722 fixes the discard; those constructors now adopt the provider on the already-foreign path. Repoint the patch at the branch carrying it and rebind both codecs onto the fork. Verified the branch carries everything already pinned rather than trusting the commit graph, which reports the two as diverged: across 3811 files the only differences are the four constructors from the fix, and `datafusion/ffi/src/session/mod.rs` is byte-identical, so the `create_physical_plan` codec fix arrives as its branch-55 backport. `ancestors` and its helpers are deleted. They existed only to keep the pre-fork session alive for a codec that could not be moved off it, and a codec bound to the running session needs no such anchor. Three tests, replacing two that were weaker than they looked. One registers a function on the fork after the codec was installed on its parent and resolves it, which is the direct evidence the rebind happened; it failed before this change. One installs a planner twice and asserts the first context still cannot resolve a function registered only on the second, covering the clone-before-adopt half — a rebind that mutated the shared handle would pass the first test and fail this one. The third keeps the live-handle case. The test it replaces required a name registered nowhere, so it passed for the same reason as the negative control and never exercised a fork at all. Note the version floor in `Cargo.toml` rather than raising it now: the patched branch still reports 55.0.0, so the requirement can only move to 55.1.0 when the patch section is removed. Building against 55.0.0 without the patch would compile and silently skip the rebind. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

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

timsaucer and others added 4 commits August 26, 2026 15:50

timsaucer added the api change label Aug 27, 2026

This was referenced Aug 27, 2026

FFI constructors silently discard arguments when the input is already foreign apache/datafusion#24722

Open

fix: preserve provided arguments during FFI object construction apache/datafusion#24723

Open

timsaucer and others added 3 commits August 27, 2026 13:34

Update temporary DF version with corrections in FFI

1e1d1aa

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

9696956

timsaucer mentioned this pull request Aug 27, 2026

FFI Query Planner on 54.0.0 #1636

Closed

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!

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