Add atomic SessionContext.with_extensions API - #1679
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>
f152621 to
725fa74
Compare
The docs build runs Sphinx with --fail-on-warning. SessionExtensionComponents documented its fields in both a napoleon `Attributes:` section and the dataclass class-body annotations, so autoapi emitted each field twice and the build failed with six "duplicate object description" warnings. Move each field's description to a per-field docstring under its annotation so autoapi renders exactly one entry per field. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
ntjohnson1
left a comment
There was a problem hiding this comment.
I don't have a fully coherent thought here. IIUC this is mostly to manage life times across the FFI boundary. I do wonder if there is a slightly cleaner way to mange this but this seems fine for now to provide a safer avenue.
I wonder if it makes sense to have a todo for some datafusion python extension skill/s. Being able to generate the 3 library example from the skill might be a nice smoke test to verify. I suspect after getting things setup for ballista/datafusion-distributed keeping it up to date shouldn't be too bad but I do suspect it will require some guidance to make sure they are doing it safely.
| extensions: One or more objects implementing | ||
| ``__datafusion_session_extension__`` (see | ||
| :py:class:`SessionExtensionExportable`). |
There was a problem hiding this comment.
This is redundant with the type hint
| supplies a query planner. | ||
|
|
||
| Examples: | ||
| >>> from my_extension import DistributedEngineExtension # doctest: +SKIP |
There was a problem hiding this comment.
It might be nice to have a little sample we import as a part of the pytest conf so these examples don't go stale. Through the stacked PRs lots of doctest skip
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.