Skip to content

Add FFI query planner support with composable extension codecs - #1672

Closed
timsaucer wants to merge 10 commits into
apache:mainfrom
timsaucer:feat/ffi-query-planner
Closed

Add FFI query planner support with composable extension codecs#1672
timsaucer wants to merge 10 commits into
apache:mainfrom
timsaucer:feat/ffi-query-planner

Conversation

@timsaucer

@timsaucer timsaucer commented Aug 7, 2026

Copy link
Copy Markdown
Member

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-distributed integration can build on.

Rationale for this change

Extension libraries (for example distributed execution engines) need to supply their own QueryPlanner to a SessionContext without compiling against the datafusion-python crate. 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).
  • A RuntimeAwareQueryPlanner adapter binds foreign planners to the Tokio runtime owned by datafusion-python.
  • New example crate datafusion-ffi-query-planner-example demonstrating a real three-library plan exchange (host, provider library, planner library as separate cdylibs), including session config transfer via SessionConfig.with_extension.

Composable extension codecs

  • with_logical_extension_codec / with_physical_extension_codec now 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.
  • 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 in the chain fails, the errors are aggregated so the owning codec's diagnostic is not masked by the default codec's generic error.
  • Fixed a latent bug where installing a codec silently reset python_udf_inlining back to enabled.

Atomic extension installation: SessionContext.with_extensions

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.

  • 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.

Documentation and tests

  • docs/source/contributor-guide/ffi.md gains 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_extensions documented as the preferred API, low-level chaining kept as advanced), including a full three-library registration recipe.
  • Rust unit tests for chain dispatch semantics; Python integration tests for codec composition on both layers plus an end-to-end three-library test with composed codecs under a foreign planner.
  • with_extensions tests 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:

  • New APIs: SessionContext.with_query_planner, SessionContext.__datafusion_query_planner__, SessionContext.with_extensions, SessionExtensionComponents, and SessionExtensionExportable.
  • Behavior change: with_logical_extension_codec / with_physical_extension_codec now compose with previously installed codecs instead of replacing them. Sessions that install a single codec are unaffected.
  • These methods no longer reset the with_python_udf_inlining setting.
  • New example package datafusion-ffi-query-planner-example in the examples folder (not shipped in the wheel).

🤖 Generated with Claude Code

timsaucer and others added 10 commits August 6, 2026 10:54
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>
@timsaucer

Copy link
Copy Markdown
Member Author

Closing in favor of a stacked split for easier review: #1677 (FFI query planner support), #1678 (composable extension codecs), #1679 (atomic SessionContext.with_extensions). Same commits, same final tree.

@timsaucer timsaucer closed this Aug 7, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant