feat: Implement FFI_QueryPlanner - #24028
Conversation
…ncryption feature is not enabled
Add FFI_QueryPlanner and ForeignQueryPlanner with logical and physical plan codec support. Forward query planners, logical optimization, and physical optimizer rules through foreign sessions, with unit and cross-library coverage.\n\nAI Disclosure: This code was written in part by an AI agent.
timsaucer
added
api change
|
|
||
| #[cfg(not(feature = "parquet_encryption"))] | ||
| #[expect(dead_code)] | ||
| #[expect(clippy::unused_async)] |
There was a problem hiding this comment.
Drive by fix to when you run clippy and don't have the parquet_encryption feature enabled.
| } | ||
|
|
||
| #[cfg(not(feature = "parquet_encryption"))] | ||
| #[expect(clippy::unused_async)] |
There was a problem hiding this comment.
Drive by fix to when you run clippy and don't have the parquet_encryption feature enabled.
| } | ||
|
|
||
| #[cfg(not(feature = "parquet_encryption"))] | ||
| #[expect(clippy::unused_async)] |
There was a problem hiding this comment.
Drive by fix to when you run clippy and don't have the parquet_encryption feature enabled.
|
Thank you for opening this pull request! Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch). Details |
Keep the standard FFI_SessionRef constructor focused on the required logical codec and derive a default physical codec. Add an explicit constructor for callers that already own matching logical and physical codecs.\n\nAI Disclosure: This code was written in part by an AI agent.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #24028 +/- ##
==========================================
+ Coverage 80.91% 80.93% +0.01%
==========================================
Files 1103 1104 +1
Lines 377219 377529 +310
Branches 377219 377529 +310
==========================================
+ Hits 305244 305555 +311
+ Misses 53775 53757 -18
- Partials 18200 18217 +17 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Document the query planner serialization boundary and exercise an A/B/C ownership model with independently loaded cdylib images. Reuse the existing FFI table provider and verify foreign plans are reconstructed as local nodes through A's codecs. AI Disclosure: This code was written in part by an AI agent.
timsaucer
changed the title
Use one test query planner for both the basic round trip and the three-library table scan scenario. Remove the redundant library C constructor from the integration-test module. AI Disclosure: This code was written in part by an AI agent.
Render the private FFI_SessionRef type as code so public query planner documentation passes rustdoc's private intra-doc link checks. AI Disclosure: This code was written in part by an AI agent.
paleolimbot
left a comment
paleolimbot
left a comment
There was a problem hiding this comment.
I will try to circle back to take a closer look at the details, but I did take two passes and didn't spot anything out of place here. You've documented the motivation nicely here and conceptually I think this is a great approach that allows quite a lot of flexibility for multiple datafusion-based libraries to interact. The FFI pattern you've established from previous PRs works well here.
I believe this will allow our (SedonaDB) spatial join to be used in datafusion-python plans (injected via a combination of logical optimizer rules and the query planner). Exciting!
milenkovicm
left a comment
milenkovicm
left a comment
There was a problem hiding this comment.
Hey @timsaucer first of all appologies for late review and thank you for sticking with this. Can't wait for this release to finally get py ballista support.
I have fiew minor questions, more for my understanding than blockers
|
|
||
| /// Library A's logical codec stores library B's provider while the logical | ||
| /// plan crosses into library C. A real application would encode enough | ||
| /// metadata to reconstruct or locate the provider instead. |
There was a problem hiding this comment.
just for my reference, which metadata is needed? will need some help to integrate this with py ballista
There was a problem hiding this comment.
I think "metadata" was poor phrasing. Basically you need enough data to faithfully reproduce the plan on both sides of the (de)serialization. I'm updating the docstring to make it more clear.
| &Self, | ||
| logical_plan_serialized: SVec<u8>, | ||
| session: FFI_SessionRef, | ||
| ) -> FfiFuture<FFI_Result<SVec<u8>>>, |
There was a problem hiding this comment.
One question, why FFI_Result<FFI_ExecutionPlan> not used instad of serialized plan?
There was a problem hiding this comment.
The reason for this is the case that is illustrated in the integration test test_three_library_query_planner_restores_type_identity.
Suppose we want one query planner to call another. Each of these planners tend to rely on downcasting the execution plans. By making the FFI query planner calls return the serialized plans, then it gets decoded within each library locally. By doing this we should get local versions of the well known execution plans that the default codec can provide.
Library C typically captures library A's query planner, then A installs C's planner on its session. C plans by delegating back to the captured handle. This is the deployment that requires serialized plans in both directions: C must downcast the nodes A produced in order to rewrite them, and A must downcast the nodes C produced to run its own passes over the result. Add an integration test for that topology. Replacing either serialization step with an FFI_ExecutionPlan handoff makes it fail, which the prior tests could not detect on the inbound leg. The test also asserts that after the swap the session reports C's own planner, documenting why C must delegate to the captured handle rather than call Session::query_planner or Session::create_physical_plan, both of which are self-references at that point. Widen the test physical codec to accept an A-local node during encode. FFI_ExecutionPlan::new unwraps a ForeignExecutionPlan back to its origin handle, so when C serializes a node A gave it, A is asked to encode the very plan its own try_decode produced. Take the delegate planner as FFI_Option<FFI_QueryPlanner> on the existing create_query_planner module entry instead of adding a second entry point. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
| physical_codec.unwrap_or_else(|| Arc::new(DefaultPhysicalExtensionCodec {})); | ||
| let physical_codec = | ||
| FFI_PhysicalExtensionCodec::new(physical_codec, runtime, task_ctx_provider); | ||
| Self::new_with_ffi_codecs(planner, logical_codec, physical_codec) |
There was a problem hiding this comment.
If None is passed to logical_codec or physical_codec, those arguments would default to DefaultLogicalExtensionCodec and DefaultPhysicalExtensionCodec, replacing any previously registered codecs even if users provide None as arguments.
🤔 there seems to be a couple of footguns with the FFI_QueryPlanner::new and FFI_QueryPlanner::new_with_ffi_codecs public APIs, maybe there's a cleaner constructor method we can expose publicly that makes these things irrepresentable?
Drop redundant `Send`/`Sync` bounds and remove two codec footguns in the FFI query planner API, per review on apache#24028. - `Session` already requires `Send + Sync`, so `&(dyn Session + Send + Sync)` was noise. Narrowing to `&dyn Session` also widens what callers can pass. - `LogicalExtensionCodec` and `PhysicalExtensionCodec` already require `Send`. The `+ Send` on the codec constructor parameters bought nothing and blocked callers holding an existing `Arc<dyn PhysicalExtensionCodec>`, since Rust will not coerce that to `Arc<dyn PhysicalExtensionCodec + Send>`. - `FFI_QueryPlanner::new_with_ffi_codecs` silently dropped the supplied codecs when re-exporting an already-foreign planner. It now adopts them while keeping the original planner identity. - `FFI_QueryPlanner::new` no longer takes `Option` codecs. Passing `None` used to install the default codecs, clobbering extension-node handling; requiring an explicit codec makes that unrepresentable. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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.: 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.:
AI Disclosure: This code was written in part by an AI agent.:
AI Disclosure: This code was written in part by an AI agent.:
github-actions
Bot
added
the
documentation
5 participants
Which issue does this PR close?
datafusion-distributedwithdatafusion-pythondatafusion-python#1612Rationale for this change
This is the last in a series of PRs that would enable FFI
Sessionto support aQueryPlanner. The prior work was inPhysicalPlanningContextexplicitly through planner traits #23649With those changes in place we now have the dependencies correct that we can expose a
FFI_QueryPlanneron aFFI_Session. With this we can enable foreign libraries such asdatafusion-distributedandballistato provide a query planner in Python and connect it directly to adatafusion-python'sSessionContext.What changes are included in this PR?
Addition only. Adds these functions to
FFI_Sessionand their supporting structures:query_planner()optimize()physical_optimizers()Are these changes tested?
Unit and integration tests are provided.
Are there any user-facing changes?
This is addition, but it does break the FFI ABI, which is already evolving in DF55.