Add atomic SessionContext.with_extensions API - #1679
Open
timsaucer wants to merge 5 commits into
Open
Conversation
Installing FFI extension codecs and query planners by chaining the existing with_* methods can bind task-context providers to intermediate contexts that are later collected, breaking the weak provider reference over the FFI boundary. with_extensions creates one destination context, passes it to each extension factory so components bind to that exact context, and installs everything in a single state write. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
MyPlannerExtension in the query-planner example crate implements the __datafusion_session_extension__ protocol from Rust: it extracts the destination context's task-context provider, binds fresh observing codecs and a planner to it, and returns SessionExtensionComponents. Its codecs record the max_rows config value resolved through the weak provider, letting tests prove the provider targets the returned context rather than the source. Documents with_extensions as the preferred API in the FFI guide. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
A DataFrame does not keep its SessionContext alive. FFI components hold a weak task-context provider, so operations that reach an FFI codec after the context is collected fail with a clean out-of-scope error rather than crashing. Lock that behavior in with a test and document the ownership contract in the FFI guide and with_extensions docstring. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Single-underscore methods on internal pyo3 classes (such as SessionContext._install_extensions) are private support methods for the Python wrappers and do not require a public wrapper. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This was referenced Aug 7, 2026
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
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?
Part 3 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
Working on the Ballista integration showed that chaining the low-level
with_*methods is easy to get wrong: FFI codecs and planners carry a weak task-context provider bound to the context they were created against, so components can end up bound to an intermediate context that is later garbage collected. Queries then fail withTaskContextProvider went out of scope over FFI boundary, or worse, silently read stale session state.What changes are included in this PR?
SessionContext.with_extensions(*extensions)installs one or more extension bundles atomically. Each bundle implements the new__datafusion_session_extension__(ctx)protocol: it receives the destination context, creates fresh components bound to that exact context, and returns them as aSessionExtensionComponents(new public dataclass;SessionExtensionExportableis the matching typing protocol).state_ref(). No context is derived after any provider is created, so every weak provider targets the returned context.DataFrameoutliving it fails with a clean out-of-scope error rather than crashing.MyPlannerExtensionin the example crate is a complete Rust implementation of the protocol, including extracting the host's task-context provider from the supplied context. Its codecs record the config value they resolve through the weak provider, letting tests prove the provider targets the returned context rather than the source.docs/source/contributor-guide/ffi.mddocumentswith_extensionsas the preferred API for extension bundles, keeps low-level chaining as advanced usage, and includes a full three-library registration recipe.Are there any user-facing changes?
New public APIs:
SessionContext.with_extensions,SessionExtensionComponents, and theSessionExtensionExportable/__datafusion_session_extension__protocol. The context-outlives-DataFrame ownership contract is now documented. No breaking changes to existing APIs.