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 1098 +9
Lines 369164 373580 +4416
Branches 369164 373580 +4416
==========================================
+ Hits 298127 301633 +3506
- Misses 53282 53954 +672
- Partials 17755 17993 +238 ☔ 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 |
| /// 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:
There was a problem hiding this comment.
Using the default physical optimizer rules would be difficult from a dependency perspective. Instead I strengthened up the warning on the trait.
| 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.
Merge the two fragmented `use datafusion_session::...` lines into one import in the external-crate group, removing the stranded entry from the `crate::` block. Addresses import-ordering review nit. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…n crate Recover the fuller rustdoc that was trimmed during the move: - PhysicalPlanner::create_physical_expr param breakdown and planning_ctx threading explanation - ExtensionPlanner::plan_table_scan return-value contract and rustdoc example (updated to &dyn Session / datafusion_session paths) - PhysicalOptimizerRule registration pointer and per-method guidance - QueryPlanner method phrasing Rendered TableProvider as plain code instead of an intra-doc link since datafusion-session does not depend on datafusion-catalog. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Make explicit that the no-op default for Session::optimize and the empty-slice default for Session::physical_optimizers are almost never correct: without them queries run unoptimized or fail (physical rules provide correctness-critical repartition/coalesce rewrites). The defaults exist only so this crate avoids an optimizer dependency; real sessions should override, e.g. by delegating to SessionState. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Make query_planner a defaulted method returning UnsupportedQueryPlanner, matching the other defaulted planning methods (optimize, physical_optimizers, statistics_registry). This keeps the planning capability a single coherent, fully-defaulted group: sessions that do not perform physical planning need no stubs at all, while sessions that drive DefaultPhysicalPlanner override the group (typically by delegating to SessionState). Removes the now-redundant UnsupportedQueryPlanner stubs from the FFI ForeignSession and the test MockSession impls, and updates the 55.0.0 upgrade guide to describe the defaulted group and the delegate-to- SessionState override pattern. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add a comment noting the inherent-vs-trait method name collision: the qualified path targets SessionState's inherent query_planner; a bare self.query_planner() would recurse into the trait method. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
I don't know why I can't respond on the thread for this comment, but my thought was to first expose what we need for the FFI_QueryPlanner right now and keep #23678 open for follow on work. |
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.