Add FFI query planner support - #1677
Open
timsaucer wants to merge 32 commits into
Open
Conversation
AI Disclosure: This code was written in part by an AI agent.:
AI Disclosure: This code was written in part by an AI agent.:
AI Disclosure: This code was written in part by an AI agent.:
timsaucer
commented
Aug 7, 2026
| Ok(()) | ||
| } | ||
|
|
||
| pub fn with_query_planner(&self, planner: Bound<'_, PyAny>) -> PyDataFusionResult<Self> { |
Member
Author
There was a problem hiding this comment.
This API is the main reason for this PR. Here we allow changing out the default query planner with a user provided query planner.
Comment on lines
+199
to
+207
| - 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 |
Member
Author
There was a problem hiding this comment.
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.
Comment on lines
+236
to
+238
| struct RuntimeAwareQueryPlanner { | ||
| planner: FFI_QueryPlanner, | ||
| } |
Member
Author
There was a problem hiding this comment.
As the docstring says, the purpose of this is to make sure we attach the runtime handle when needed.
Comment on lines
+1456
to
+1459
| pub fn __datafusion_query_planner__<'py>( | ||
| &self, | ||
| py: Python<'py>, | ||
| ) -> PyResult<Bound<'py, PyCapsule>> { |
Member
Author
There was a problem hiding this comment.
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.
Comment on lines
+35
to
+38
| #[derive(Clone, Debug)] | ||
| pub(crate) struct PlannerConfig { | ||
| pub max_rows: usize, | ||
| } |
Member
Author
There was a problem hiding this comment.
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.
Contributor
There was a problem hiding this comment.
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
approved these changes
Aug 14, 2026
ntjohnson1
left a comment
ntjohnson1
left a comment
Contributor
There was a problem hiding this comment.
Appears consistent with the rest of the FFI plumbing
| """ | ||
| self.ctx.add_physical_optimizer_rule(rule) | ||
|
|
||
| def with_query_planner( |
Contributor
There was a problem hiding this comment.
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)
Member
Author
There was a problem hiding this comment.
Good point! Also worth updating the skill to match this pattern
milenkovicm
approved these changes
Aug 17, 2026
milenkovicm
left a comment
milenkovicm
left a comment
Contributor
There was a problem hiding this comment.
thanks @timsaucer cant want to get this integrated
| } | ||
|
|
||
| #[pymethods] | ||
| impl PlannerConfig { |
Contributor
There was a problem hiding this comment.
Nit, MyPlannerConfig to have names aligned,
Comment on lines
+35
to
+38
| #[derive(Clone, Debug)] | ||
| pub(crate) struct PlannerConfig { | ||
| pub max_rows: usize, | ||
| } |
Contributor
There was a problem hiding this comment.
this is needed for ballista, thanks Tim for example
| observations: Arc::clone(&self.observations), | ||
| }); | ||
| let runtime = get_tokio_runtime().handle().clone(); | ||
| let ctx_provider = Arc::new(SessionContext::new()) as Arc<dyn TaskContextProvider>; |
Contributor
There was a problem hiding this comment.
is this session context be parameter of method call on the line 119 ? are those two different sessions ?
Member
Author
There was a problem hiding this comment.
Really good catch! This led me down a rabbit hole and I ended up needing two upstream fixes:
Session::create_physical_planover FFI ignores the session'sLogicalExtensionCodecdatafusion#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>
`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>
This was referenced Aug 27, 2026
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>
Regenerating the lock against the patched DataFusion fork silently downgraded base64 from 0.23.1 to 0.23.0. Nothing requires the older version -- neither the fork nor upstream 55.0.0 constrains it -- so this was incidental churn from the lockfile refresh, not a resolution result. Restores the checksum main already had and re-points the three dependents (datafusion-common, datafusion-functions, parquet). No other dependency moves; cargo metadata --locked still resolves cleanly. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Installing a foreign query planner forks the session state, and the fork was minting a new session id. SessionStateBuilder::new_from_existing drops the id and build() replaces it with a fresh UUID, while SessionContext had already cached the original into a field of its own back at new_with_state. Overwriting the state in place afterwards left the two disagreeing: session_id() returned the pre-fork id, every TaskContext handed to a foreign codec carried a different one. Nothing in DataFusion core keys on the session id beyond debug logging, so this broke no in-tree behavior. It matters at the FFI boundary, where session id equality is the idiom for "which session is this codec bound to", and for extension libraries correlating host-side and worker-side state. Upstream hit the same case in SessionContext::enable_url_table and preserves the id explicitly, guarded by preserve_session_context_id. Passing the id through the builder makes the fork, its state, and its TaskContexts agree, which is what the derived_parts doc comment and the FFI contributor guide already claimed. Verified by reading the id out of a decode callback via the example codec's require_udf_on_decode error path -- the only way to observe the state-side id from Python -- with and without the fix. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Same drift just fixed in derived_parts, but on a path that never forks. add_physical_optimizer_rule rebuilds SessionState through SessionStateBuilder::new_from_existing and writes it straight back into the caller's own session, so the fresh id build() mints replaces the one SessionContext had already cached at construction. The session the user is holding then reports one id from session_id() and a different one from every TaskContext it hands out, with no derivation to explain it. Reproduced against a foreign codec, reading the id back out of a decode callback: identical setup differing only by an add_physical_optimizer_rule call went from MATCH to DRIFT, and back to MATCH with the id threaded through the builder. This is the last new_from_existing call site in the crate. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The two session id fixes had no regression guard. A Python-level assertion cannot provide one: session_id() reads a copy SessionContext caches at construction, which stayed correct through both bugs. The id that actually moved was the one inside the TaskContext handed to a foreign codec's decode callback, which nothing exposed. Give the example codecs a TaskContextProbe that records it. This replaces the bare AtomicUsize the require_udf_on_decode support used, so the counter and the session id are recorded together, and the id is recorded on every decode rather than only when a function was requested. Three tests, all against the codec-side id rather than session_id(): a fork agrees with its codecs, add_physical_optimizer_rule does not move the id, and a two-deep fork chain leaves both halves on the parent's id. Confirmed non-vacuous: with both fixes reverted all three fail and the other 17 tests pass; with only the derived_parts fix restored, exactly the add_physical_optimizer_rule test still fails. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
call_capsule_getter rewrote every TypeError from a capsule getter into "Incompatible libraries ... Upgrade the library providing this object", and dropped the original. Only an arity mismatch means the library is out of date. An extension author whose own getter raised a TypeError -- a bad cast, a wrong argument to something it called -- was told the error was a version problem and lost the error that would have located it. The two are distinguishable without guessing at message text: an arity mismatch is raised by the call machinery before the getter's frame exists, so no frame unwinds and no traceback is attached, while an error from the body carries one. Verified to hold for both pure-Python and pyo3-compiled getters, which is the case that matters here since extension libraries are compiled. Also chains the original as __cause__ on the paths that do report an upgrade, so the arity error stays readable. Tests cover all three outcomes. Confirmed non-vacuous: dropping the traceback check fails only the inside-the-getter test, dropping set_cause fails only the upgrade test. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The from_pycapsule! macros call the getter with no arguments. That is correct for __datafusion_physical_optimizer_rule__ and __datafusion_task_context_provider__, which take no session, but __datafusion_physical_extension_codec__ now takes the session it is being installed on, so this helper was the one member of the family left speaking the old protocol. Nothing in the tree called it, but datafusion-python-util is published by `cargo publish --workspace`, so it was still reachable. Against an updated codec it raised a bare TypeError, bypassing the ImportError that names the method. Against an outdated one it succeeded and produced a codec resolving names against the wrong session -- the silent failure the rest of this work exists to prevent. Removing it is a breaking change to that crate, but the crate already breaks this release: ffi_logical_codec_from_pycapsule gained its session parameter. A compile error pointing at the replacement beats a helper that quietly binds to nothing. Callers move to ffi_physical_codec_from_pycapsule, which passes the session, plus (&ffi).into() where an Arc<dyn PhysicalExtensionCodec> is wanted -- what crates/core already does. Documents both helper changes in the 55.0.0 upgrade guide, which until now covered only the __datafusion_*__ method signatures and not the Rust helpers the same authors call. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Only ffi_query_planner_from_pycapsule validated the version a capsule
reported. The codec and table provider importers dereference a foreign
struct through the same `unsafe { data.as_ref() }` and were happy to
accept one built against a different DataFusion.
Extracts the planner's inline check into check_ffi_version and applies
it to the logical codec, physical codec, and table provider importers as
well. The helper is pub so extension libraries writing their own
importers can use it.
Two things the symmetry cannot reach, both now documented where someone
would look:
FFI_TaskContextProvider, FFI_TableProviderFactory, and
FFI_ExtensionOptions carry no version field, so their importers cannot
check. The from_pycapsule!/try_from_pycapsule! macros are #[macro_export]
and generic over the FFI type, so requiring a version field there would
break downstream users holding one of those three; they stay unchecked
and their doc comment now says to call check_ffi_version directly.
This is a diagnostic, not a soundness guarantee, and the helper says so:
`version` is not the first field on any of these structs, so reading it
already assumes the local layout. It turns the realistic failure -- a
library compiled against a different DataFusion -- into a clear error
instead of undefined behaviour on first use, which is what
datafusion_ffi::version is documented to be for.
Verified all four sites are wired by inverting the comparison and
confirming each one fires from the test suites.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Exact equality is only right while datafusion_ffi::version tracks the crate's semver major, which it does today, so the number moves on every major release whether or not the ABI changed. If a version span later becomes compatible, a maintainer needs to know that this one body holds the whole policy -- callers pass a value and no decision -- and that relaxing it at a call site would reintroduce the split the helper was added to remove. Also records the likelier resolution: if the ABI is stable but version still follows the crate major, upstream's compatibility marker is wrong for every consumer, so the fix belongs there rather than in a local range policy. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The table provider and table function importers each carried their own copy of the TypeError-to-ImportError mapping, predating call_capsule_getter and never folded into it. Both therefore missed the correction it since received: they rewrote a TypeError raised inside a correctly-signed getter into "upgrade your library", and discarded the original. Three copies of one mapping, two of them stale, is the reason to have one. Both now call the shared helper, so they pick up the traceback discrimination and the __cause__ chain, and any later correction reaches all three by construction. Their messages named DataFusion 52.0.0. The shared message names the method that refused the argument instead, which points at the specific hook rather than a release, and the upgrade guide carries the version detail. call_capsule_getter is now pub, with a doc comment saying to use it rather than calling getattr directly. Tests cover both outcomes on both paths. Verified against the previous build that they are non-vacuous: before this change the raises-inside case produced the same misleading ImportError as the old-signature case. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The before and after snippets pass the task context provider differently, by reference in one and by value in the other, with nothing saying why. Read as a diff it looks like a typo in one of them, and a reader correcting it would be puzzled when both versions compile. Both are valid: the parameter is impl Into<FFI_TaskContextProvider>, which is satisfied by &Arc<dyn TaskContextProvider> and by FFI_TaskContextProvider itself, and the latter is what ffi_task_context_provider_from_pycapsule returns. The argument changes because the provider now comes from the session instead of a field, which is the point of the migration. The contributor guide shows only the post-migration form, so it needs no equivalent note. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The README already describes these as one-shot registries that consume each token during decoding, but the source comments did not, and the source is what someone reuses the pattern from. The existing comment warned that the registry is process-local without saying that a decode removes its entry, which is the constraint most likely to bite. Documents both consequences on the registry accessors, where the mechanism lives, with a pointer from each struct doc: - Decode consumes the token, so the same encoded bytes cannot be decoded twice. Fine here because every plan is encoded immediately before the one decode that consumes it, but it rules out replaying a stored plan, retrying a decode, or fanning one plan out to several readers. - An encode that never reaches a decoder leaks for the life of the process. Normal operation does not: encode and decode counts balance exactly across repeated queries, which is what makes remove-on-decode the right trade here rather than a leak on every call. Comments only. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
MyQueryPlanner::new imported its fallback immediately, with no session to pass, so the fallback's getter was called with no arguments. That works for a SessionContext, whose getter takes the session optionally, and for a raw capsule, which has no getter at all. It fails for another foreign planner, which implements the same protocol this type does and requires the argument -- and layering on another planner is the case a distributed engine actually needs. The docstring claimed fallback "takes anything exporting __datafusion_query_planner__", which was not true. Holds the Python object instead and imports it in __datafusion_query_planner__, where the session is in hand and can be forwarded. All three fallback kinds now work. Deferring also removes a footgun rather than adding one. Passing a SessionContext now delegates to whichever planner it holds at install time, and since with_query_planner calls the getter before installing, the context still reports its previous planner, so wrapping a context in a planner installed on that same context does not recurse. Arc<Py<PyAny>> rather than Py<PyAny> because pyo3 0.29 gates Py: Clone behind the py-clone feature, and this type derives Clone. Matches how PythonTableFunctionCallable holds its callable. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
foreign_session, foreign_provider, and foreign_plan were written with store, so each one described only the most recent plan. Their accessors are named foreign_*_observed, which asks whether the thing was ever seen, and the tests assert them after running more than one query. The existing tests passed by luck. Reproduced: after scanning a foreign provider and then running SELECT 1, foreign_provider_observed goes from True back to False. Writes them with fetch_or so a later plan cannot retract what an earlier one observed. plan_calls already accumulated, used_fallback only ever stores true so it was already cumulative, and last_max_rows is deliberately last-wins as its name says. Documents that split on the struct, since it is the kind of thing that gets "tidied" back. Confirmed non-vacuous: with store restored, exactly the new test fails and the other 22 pass. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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-distributedintegration 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
QueryPlannerto aSessionContextwithout compiling against thedatafusion-pythoncrate. 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).__datafusion_query_planner__,__datafusion_logical_extension_codec__, and__datafusion_physical_extension_codec__— now receive theSessionContextthey are being installed on, matching what__datafusion_table_provider__and friends have done since 52.0.0. A codec takes theTaskContextProviderfrom it; a planner takes both codecs. Neither example constructs aSessionContextany more, and decode callbacks now resolve names against the session running the query.PySessionContextretains forked ancestors, so a codec bound to a session before a planner fork cannot be left holding a dropped weak reference.datafusion-ffi-query-planner-exampledemonstrating a real three-library plan exchange (host, provider library, planner library as separate cdylibs), including session config transfer viaSessionConfig.with_extension.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.docs/source/contributor-guide/ffi.mdsections 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 fromAGENTS.md. Correcteduser-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_plannerandSessionContext.__datafusion_query_planner__.Breaking:
__datafusion_logical_extension_codec__and__datafusion_physical_extension_codec__now take asession: Bound<PyAny>parameter, so any extension library implementing them must be updated.docs/source/user-guide/upgrade-guides.mdhas a 55.0.0 section with before and after. Calling the old signature raises an import error naming the method rather than a bareTypeError.SessionContext's own getters accept the argument optionally, soctx.__datafusion_logical_extension_codec__()is unaffected.A new example crate ships under
examples/.