Skip to content

Add atomic SessionContext.with_extensions API - #1679

Open
timsaucer wants to merge 5 commits into
feat/ffi-composable-codecsfrom
feat/ffi-with-extensions
Open

Add atomic SessionContext.with_extensions API#1679
timsaucer wants to merge 5 commits into
feat/ffi-composable-codecsfrom
feat/ffi-with-extensions

Conversation

@timsaucer

@timsaucer timsaucer commented Aug 7, 2026

Copy link
Copy Markdown
Member

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 with TaskContextProvider 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 a SessionExtensionComponents (new public dataclass; SessionExtensionExportable is the matching typing protocol).
  • Internally one destination context is created, factories bind against it, all codecs are installed, and the planner is bound against the final codec chains in a single state write through that context's own state_ref(). No context is derived after any provider is created, so every weak provider targets the returned context.
  • Validation happens before any state change: a failing factory or invalid capsule leaves the source context untouched.
  • If no bundle supplies a planner, an existing FFI planner on the source context is rebound to the new codec chains; at most one bundle may supply a planner.
  • The ownership contract is documented and tested: the returned context is the strong owner, and a DataFrame outliving it fails with a clean out-of-scope error rather than crashing.
  • MyPlannerExtension in 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.md documents with_extensions as 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 the SessionExtensionExportable / __datafusion_session_extension__ protocol. The context-outlives-DataFrame ownership contract is now documented. No breaking changes to existing APIs.

timsaucer and others added 5 commits August 7, 2026 15:19
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant