Add FFI query planner support with composable extension codecs - #1672
Closed
timsaucer wants to merge 10 commits into
Closed
Add FFI query planner support with composable extension codecs#1672timsaucer wants to merge 10 commits into
timsaucer wants to merge 10 commits into
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.:
Installing a logical or physical extension codec now prepends it to a codec chain instead of replacing the prior codec. The most recently installed codec is consulted first, falling through codec by codec to the default codec. This lets multiple independent extension libraries install codecs on the same session, and removes the codec registration ordering requirement between libraries. Chain dispatch treats a codec error as "not mine". Encoding runs each codec against a scratch buffer so failed attempts leave no partial bytes, and treats Ok-with-no-bytes (encode by name) as no opinion so later codecs still get a chance. When every codec fails, the errors are aggregated so the owning codec's diagnostic is not masked by the default codec's generic error. Also preserves the python_udf_inlining setting when installing a codec; previously it was silently reset to enabled. Documents the remaining planner constraint: a session holds one query planner, layering is explicit via fallback capsules, and codecs must be installed before exporting or chaining planners because a planner capsule captures the codecs at export time. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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>
Member
Author
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?
Related to #1612. This PR does not close it, but provides the FFI query planner and codec plumbing that a
datafusion-distributedintegration can build on.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.Supporting a foreign planner also surfaced a codec problem: a query can involve three independent native libraries (datafusion-python, a provider library, and a planner library), and each library needs its extension codecs active on the session at the same time. Previously, installing a logical or physical extension codec replaced the prior codec, so the second library's install silently discarded the first — plans then failed later with a confusing decode error. Codecs now compose.
What changes are included in this PR?
FFI query planner
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).RuntimeAwareQueryPlanneradapter binds foreign planners to the Tokio runtime owned by datafusion-python.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.Composable extension codecs
with_logical_extension_codec/with_physical_extension_codecnow prepend to a codec chain instead of replacing the prior codec. The most recently installed codec is consulted first, falling through codec by codec to DataFusion's default codec. A codec signals "not mine" by returning an error.python_udf_inliningback to enabled.Atomic extension installation:
SessionContext.with_extensionsWorking 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.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.Documentation and tests
docs/source/contributor-guide/ffi.mdgains sections on composable codecs (family-prefix discipline, registration order between libraries no longer matters), planner layering (install all codecs before exporting or chaining planners, since a planner capsule captures the codecs at export time), and extension bundles (with_extensionsdocumented as the preferred API, low-level chaining kept as advanced), including a full three-library registration recipe.with_extensionstests cover validation, codec precedence between bundles, planner rebinding, identity (two live contexts with diverged config), GC of the source context and bundle objects, and failure atomicity.Are there any user-facing changes?
Yes:
SessionContext.with_query_planner,SessionContext.__datafusion_query_planner__,SessionContext.with_extensions,SessionExtensionComponents, andSessionExtensionExportable.with_logical_extension_codec/with_physical_extension_codecnow compose with previously installed codecs instead of replacing them. Sessions that install a single codec are unaffected.with_python_udf_inliningsetting.datafusion-ffi-query-planner-examplein the examples folder (not shipped in the wheel).🤖 Generated with Claude Code