Add FFI query planner support - #1677
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.:
| Ok(()) | ||
| } | ||
|
|
||
| pub fn with_query_planner(&self, planner: Bound<'_, PyAny>) -> PyDataFusionResult<Self> { |
There was a problem hiding this comment.
This API is the main reason for this PR. Here we allow changing out the default query planner with a user provided query planner.
| - name: Build FFI query planner test library | ||
| if: matrix.python-tag == 'abi3' | ||
| uses: PyO3/maturin-action@v1 | ||
| with: | ||
| target: x86_64-unknown-linux-gnu | ||
| manylinux: "2_28" | ||
| working-directory: examples/datafusion-ffi-query-planner-example | ||
| args: --out dist | ||
| rustup-components: rust-std |
There was a problem hiding this comment.
In order to prove that the 3 library approach works where we have different codecs and different execution plans provided, we are adding a second test library. This way we can make sure there is no accidental ability to reach into a foreign code block.
| struct RuntimeAwareQueryPlanner { | ||
| planner: FFI_QueryPlanner, | ||
| } |
There was a problem hiding this comment.
As the docstring says, the purpose of this is to make sure we attach the runtime handle when needed.
| pub fn __datafusion_query_planner__<'py>( | ||
| &self, | ||
| py: Python<'py>, | ||
| ) -> PyResult<Bound<'py, PyCapsule>> { |
There was a problem hiding this comment.
We need our session context to export it's own query planner because we have a use case where one query planner can depend on another. This is already supported by datafusion-distributed, so we want to be certain we support it here.
| #[derive(Clone, Debug)] | ||
| pub(crate) struct PlannerConfig { | ||
| pub max_rows: usize, | ||
| } |
There was a problem hiding this comment.
I'm adding this to the query planner example because it's a very common pattern that we will need custom configs for the query planner, so it is reasonable to need insurance that configs pass over the FFI boundary properly and to use as a demonstration to anyone who is providing such a library.
1 participant
Which issue does this PR close?
Related to #1612. This PR does not close it, but provides the FFI query planner plumbing that a
datafusion-distributedintegration can build on.This is part 1 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
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.What changes are included in this PR?
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.docs/source/contributor-guide/ffi.mdsections covering the query planner capsule protocol and the three-library setup.Are there any user-facing changes?
New public APIs:
SessionContext.with_query_plannerandSessionContext.__datafusion_query_planner__. A new example crate ships underexamples/. No breaking changes to existing APIs.