refactor: move planning APIs to session crate - #23842
Conversation
Move planner and physical optimizer interfaces into datafusion-session, convert planner session arguments to &dyn Session, and preserve existing public reexports.\n\nAI Disclosure: This code was written in part by an AI agent.
github-actions
Bot
added
documentation
timsaucer
changed the title
|
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 |
Make query_planner a required Session method and provide a public UnsupportedQueryPlanner for sessions that do not expose one, including ForeignSession until planner FFI support is added.\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 #23842 +/- ##
==========================================
- Coverage 80.75% 80.74% -0.02%
==========================================
Files 1089 1091 +2
Lines 369164 369340 +176
Branches 369164 369340 +176
==========================================
+ Hits 298127 298223 +96
- Misses 53282 53356 +74
- Partials 17755 17761 +6 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Move test-only imports into their test modules instead of conditionally importing them into the parent modules.\n\nAI Disclosure: This code was written in part by an AI agent.
| fn query_planner(&self) -> Arc<dyn QueryPlanner + Send + Sync> { | ||
| Arc::new(UnsupportedQueryPlanner) | ||
| } | ||
|
|
There was a problem hiding this comment.
I will open a follow up PR immediately after this merges that will replace this with an FFI equivalent implementation.
| /// Return the query planner for this session. | ||
| fn query_planner(&self) -> Arc<dyn QueryPlanner + Send + Sync>; | ||
|
|
||
| /// Optimize a logical plan. | ||
| /// | ||
| /// The default implementation returns the plan unchanged. Sessions that use | ||
| /// the default query planning implementation should override this method. | ||
| fn optimize(&self, plan: &LogicalPlan) -> Result<LogicalPlan> { | ||
| Ok(plan.clone()) | ||
| } | ||
|
|
||
| /// Return the physical optimizer rules for this session. | ||
| /// | ||
| /// The default implementation returns no rules. Sessions that use the | ||
| /// default physical planner should override this method. | ||
| fn physical_optimizers(&self) -> &[Arc<dyn PhysicalOptimizerRule + Send + Sync>] { | ||
| &[] | ||
| } | ||
|
|
||
| /// Return the optional statistics registry used during physical optimization. | ||
| fn statistics_registry(&self) -> Option<&StatisticsRegistry> { | ||
| None | ||
| } |
There was a problem hiding this comment.
These are the new methods hoisted up to Session from SessionState
There was a problem hiding this comment.
Is there anything that we should be doing for the other methods from SessionState? (expr_planners, relation_planners, table_factories, serializer_registry or function_factory).
I imagine that if any QueryPlanner or ExtensionPlanner was relying on those, they will no longer be able to use them without downcasting.
|
@milenkovicm this PR will unblock FFI_QueryPlanner ! |
There was a problem hiding this comment.
Pull request overview
This PR continues the ongoing refactor to make DataFusion’s planning and optimization extension points accessible via &dyn Session (including across FFI boundaries) by moving the relevant traits into the datafusion-session crate and preserving prior public paths via re-exports.
Changes:
- Move planner contracts (
QueryPlanner,PhysicalPlanner,ExtensionPlanner) and physical optimizer contracts (PhysicalOptimizerRule,PhysicalOptimizerContext) intodatafusion-session, and update signatures to accept&dyn Session. - Extend the
Sessiontrait withquery_plannerand add defaulted hooks (optimize,physical_optimizers,statistics_registry) to expose planning/optimization behavior without downcasting. - Update core planning/optimization wiring, FFI
ForeignSession, docs, tests, and examples; keep prior import paths working via re-exports.
Reviewed changes
Copilot reviewed 17 out of 18 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| docs/source/library-user-guide/upgrading/55.0.0.md | Documents the trait moves, signature changes (&SessionState → &dyn Session), re-exported paths, and new required Session::query_planner. |
| datafusion/session/src/session.rs | Extends Session with query_planner plus defaulted optimization/optimizer/stats-registry hooks. |
| datafusion/session/src/planner.rs | Introduces the planner trait definitions in datafusion-session, including UnsupportedQueryPlanner. |
| datafusion/session/src/physical_optimizer.rs | Introduces physical optimizer traits in datafusion-session with optional statistics-registry support. |
| datafusion/session/src/lib.rs | Exposes new planner and physical_optimizer modules and re-exports their public contracts. |
| datafusion/physical-optimizer/src/optimizer.rs | Switches optimizer traits to re-export from datafusion-session for backward compatibility. |
| datafusion/physical-optimizer/Cargo.toml | Adds datafusion-session dependency needed for the new re-export source. |
| datafusion/ffi/src/session/mod.rs | Implements Session::query_planner for ForeignSession (currently UnsupportedQueryPlanner) and adds unit coverage for new default hooks. |
| datafusion/datasource/src/url.rs | Updates test-only Session impls to satisfy new required query_planner. |
| datafusion/datasource-arrow/src/file_format.rs | Updates test-only Session impls to satisfy new required query_planner. |
| datafusion/core/tests/user_defined/user_defined_plan.rs | Updates custom planner/extension planner signatures to take &dyn Session. |
| datafusion/core/src/physical_planner.rs | Re-exports planner traits from datafusion-session, updates planner/optimizer interfaces to use &dyn Session, and adapts optimizer context wiring. |
| datafusion/core/src/execution/session_state.rs | Implements the expanded Session trait for SessionState and updates DefaultQueryPlanner signature to &dyn Session. |
| datafusion/core/src/execution/context/mod.rs | Re-exports QueryPlanner from datafusion-session to preserve prior paths. |
| datafusion/core/src/datasource/listing_table_factory.rs | Updates test-only Session impls to satisfy new required query_planner. |
| datafusion-examples/examples/relation_planner/table_sample.rs | Updates example custom planner/extension planner signatures to &dyn Session. |
| datafusion-examples/examples/dataframe/cache_factory.rs | Updates example custom planner/extension planner signatures to &dyn Session. |
| Cargo.lock | Records the updated dependency graph (physical-optimizer now depends on datafusion-session). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
Would any of you be willing to review this @andygrove @milenkovicm @gabotechs ? This hoists the caveat: distributed is still blocked by extension data on configs, but maybe there is something I can do on it. |
|
I can take a look next week!
This was addressed on datafusion-contrib/datafusion-distributed#550, let me know if that's not enough. |
No, the issue is that only string key/value pairs can make it through FFI configs. I'm going to investigate to see if the python branch I have for datafusion-distributed can pass the internal config around in a different way other than getting them out of the session's config. |
I wonder if we should make a datafusion-api crate, or datafusion-traits crate 🤔 |
Why would those extensions need to cross any FFI boundary? I'd imagine that in the Python API, we would build something like a |
| use datafusion_session::{PhysicalOptimizerContext, Session}; | ||
|
|
||
| use arrow::array::{RecordBatch, builder::StringBuilder}; | ||
| use arrow::compute::SortOptions; |
There was a problem hiding this comment.
nit: just some import ordering hygiene:
| use datafusion_session::{PhysicalOptimizerContext, Session}; | |
| use arrow::array::{RecordBatch, builder::StringBuilder}; | |
| use arrow::compute::SortOptions; | |
| use datafusion_session::{PhysicalOptimizerContext, Session}; | |
| use arrow::array::{RecordBatch, builder::StringBuilder}; |
| /// Transforms one [`ExecutionPlan`] into another that computes the same results, | ||
| /// but may do so more efficiently. | ||
| pub trait PhysicalOptimizerRule: Debug + std::any::Any { |
There was a problem hiding this comment.
Seems like we've lost some documentation while moving this here? I'd say the docs in the previous struct are still relevant no? For example:
/// `PhysicalOptimizerRule` transforms one ['ExecutionPlan'] into another which
/// computes the same results, but in a potentially more efficient way.
///
/// Use [`SessionState::add_physical_optimizer_rule`] to register additional
/// `PhysicalOptimizerRule`s.
///
/// [`SessionState::add_physical_optimizer_rule`]: https://docs.rs/datafusion/latest/datafusion/execution/session_state/struct.SessionState.html#method.add_physical_optimizer_rule
| /// The default implementation returns no rules. Sessions that use the | ||
| /// default physical planner should override this method. | ||
| fn physical_optimizers(&self) -> &[Arc<dyn PhysicalOptimizerRule + Send + Sync>] { | ||
| &[] |
There was a problem hiding this comment.
🤔 I wonder if this should be defaulting to the default physical optimizer rules.
I'm not sure if DataFusion is even useful without those:
| struct SessionOptimizerContext<'a> { | ||
| session: &'a dyn Session, | ||
| } | ||
|
|
There was a problem hiding this comment.
It's a shame that we cannot implement PhysicalOptimizerContext automatically for any dyn Session
| fn config_options(&self) -> &ConfigOptions; | ||
|
|
||
| /// Returns the statistics registry for enhanced statistics lookup. | ||
| /// | ||
| /// Returns `None` if no registry is configured, in which case rules | ||
| /// should fall back to using [`ExecutionPlan::partition_statistics`]. | ||
| fn statistics_registry(&self) -> Option<&StatisticsRegistry> { | ||
| None | ||
| } | ||
| } |
There was a problem hiding this comment.
🤔 these two methods also exist on Session. I wonder if that can introduce some ambiguity that the compiler surfaces as "multiple applicable items in scope".
An alternative that comes to mind is to rely on trait composition: instead of duplicating methods in config_options and statistics_registry. This probably applies to more traits rather than just this. Trait upcasting was stabilized a while ago in the Rust compiler.
Just throwing the idea, don't have a strong opinion here.
| /// Return the query planner for this session. | ||
| fn query_planner(&self) -> Arc<dyn QueryPlanner + Send + Sync>; | ||
|
|
||
| /// Optimize a logical plan. | ||
| /// | ||
| /// The default implementation returns the plan unchanged. Sessions that use | ||
| /// the default query planning implementation should override this method. | ||
| fn optimize(&self, plan: &LogicalPlan) -> Result<LogicalPlan> { | ||
| Ok(plan.clone()) | ||
| } | ||
|
|
||
| /// Return the physical optimizer rules for this session. | ||
| /// | ||
| /// The default implementation returns no rules. Sessions that use the | ||
| /// default physical planner should override this method. | ||
| fn physical_optimizers(&self) -> &[Arc<dyn PhysicalOptimizerRule + Send + Sync>] { | ||
| &[] | ||
| } | ||
|
|
||
| /// Return the optional statistics registry used during physical optimization. | ||
| fn statistics_registry(&self) -> Option<&StatisticsRegistry> { | ||
| None | ||
| } |
There was a problem hiding this comment.
Is there anything that we should be doing for the other methods from SessionState? (expr_planners, relation_planners, table_factories, serializer_registry or function_factory).
I imagine that if any QueryPlanner or ExtensionPlanner was relying on those, they will no longer be able to use them without downcasting.
Which issue does this PR close?
Rationale for this change
This PR unlocks using a query planner across the FFI boundary.
What changes are included in this PR?
Moves these traits to the
datafusion-sessioncrate:QueryPlannerPhysicalPlannerExtensionPlannerPhysicalOptimizerRulePhysicalOptimizerContextAdditionally changed the method signatures from taking
&SessionStateto&dyn Session.Adds these methods on
Sessiontrait:fn query_planner(&self) -> Arc<dyn QueryPlanner + Send + Sync>fn optimize(&self, plan: &LogicalPlan) -> Result<LogicalPlan>fn physical_optimizers(&self) -> &[Arc<dyn PhysicalOptimizerRule + Send + Sync>]fn statistics_registry(&self) -> Option<&StatisticsRegistry>Are these changes tested?
Unit tests are added for the new methods.
Are there any user-facing changes?
Yes, users must implement new methods for
query_plannerfor their customSession. This is not expected to impact many users, since it is most common to use the existingSessionState.