refactor: thread SubqueryContext explicitly through physical planning#22340
refactor: thread SubqueryContext explicitly through physical planning#22340timsaucer wants to merge 2 commits into
Conversation
Removes the SessionState.clone() + execution_props_mut() trick used by DefaultPhysicalPlanner to register scalar-subquery state for expression lowering. That side channel relied on stashing per-plan state in ExecutionProps, which forced physical planning to hold a mutable SessionState and blocked moving QueryPlanner / PhysicalPlanner to &dyn Session. Introduces SubqueryContext in datafusion-expr and threads it explicitly through create_initial_plan_inner, task_helper, map_logical_node_to_physical, and the standalone planning helpers. Adds *_with_subquery_context dual entry points for create_physical_expr, create_physical_exprs, create_physical_sort_expr(s), create_window_expr(_with_name), and LoweredAggregateBuilder.with_subquery_context — original public signatures are preserved and delegate with SubqueryContext::default(), so downstream callers are unaffected. Drops the unreleased subquery_indexes and subquery_results fields from ExecutionProps. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
|
@neilconway Would you mind taking a look? It's a fair amount of plumbing and I'm open to other suggestions. This is to unblock some follow on work where we want to use |
|
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 |
|
@milenkovicm this is the PR I was talking about to unblock FFI query planner |
|
Sorry @timsaucer had no time for reviews last week, will have a look when back from vacation |
# Conflicts: # datafusion/core/src/physical_planner.rs # datafusion/physical-expr/src/lib.rs # datafusion/physical-expr/src/physical_expr.rs
| pub fn create_physical_sort_exprs_with_subquery_context( | ||
| exprs: &[SortExpr], | ||
| input_dfschema: &DFSchema, | ||
| execution_props: &ExecutionProps, | ||
| subquery_ctx: &SubqueryContext, | ||
| ) -> Result<Vec<PhysicalSortExpr>> { | ||
| exprs |
There was a problem hiding this comment.
It's not very beautiful to have this two variants of the same function with different argument count and different name be part of the public API.
It does not scale very well in case newer params need to be added.
I do not have a better suggestion though...
| let mut owned = session_state.clone(); | ||
| owned.execution_props_mut().subquery_indexes = index_map; | ||
| owned.execution_props_mut().subquery_results = results.clone(); | ||
| let session_state = Cow::Owned(owned); | ||
| let subquery_ctx = SubqueryContext::new(index_map, results.clone()); |
There was a problem hiding this comment.
Given that the dyn Session trait has the .execution_props() method, and that ExecutionProps implements Clone, shouldn't it be possible to do something like this?
let mut props = session_state.execution_props().clone();
props.subquery_indexes = index_map;
props.subquery_results = results.clone();And thread around explicitly that cloned+mutated ExecutionProps everywhere?
There was a problem hiding this comment.
I gave it a try in a draft:
And it looks dead simple, but not sure if I'm missing something.
Which issue does this PR close?
None, but it is a precursor to making
QueryPlanner/PhysicalPlannertake&dyn Sessioninstead of&SessionState, so we can expose the query planner to FFI properly (similar to #22151).Rationale for this change
"Expression lowering" here means turning a logical
Exprinto a physicalArc<dyn PhysicalExpr>— the job ofcreate_physical_exprinphysical-expr/src/planner.rs. That function only receives&ExecutionProps, so any state it needs has to arrive throughExecutionProps.Uncorrelated scalar subqueries need such state: the
Expr::ScalarSubquerybranch ofcreate_physical_exprhas to look up which result slot a given subquery maps to. TodayDefaultPhysicalPlanner::create_initial_plansupplies that by cloning theSessionStateand mutating the embeddedExecutionProps:The block comment at that write site already flags this as a hack that should live in a dedicated planning context. It matters beyond aesthetics because this exact clone-and-mutate pattern is impossible once the planner takes
&dyn Session:Sessionis a trait object with noclone(), sosession_state.clone()has no equivalent — you cannot cheaply produce an owned copy to mutate.Session::execution_props(&self)returns&ExecutionPropsimmutably by design (it is called from&selfTableProviderhot paths), so there is noexecution_props_mut()to write into.So the subquery state cannot ride inside
ExecutionPropsunder&dyn Session. It has to be threaded explicitly instead. This PR does that threading now, decoupled from the larger&dyn Sessionchange.What changes are included in this PR?
SubqueryContextindatafusion-expr— a small carrier bundling theSubquery -> SubqueryIndexmap and the sharedScalarSubqueryResults.&SubqueryContextexplicitly fromDefaultPhysicalPlannerdown into expression lowering, replacing theExecutionPropsside channel.SessionState::clone()+execution_props_mut()mutation at the write site.subquery_indexesandsubquery_resultsfields fromExecutionProps(breaking — see user-facing changes).Why add
_with_subquery_contextvariants instead of changing the existing signatures?Adding a 4th
subquery_ctxparameter tocreate_physical_expris a cleaner end state but has the wrong cost:create_physical_exprhas ~92 callers in this repo plus downstream crates (datafusion-comet, ballista, datafusion-python, custom planners). All would have to pass the new argument.Expr::ScalarSubquerybranch — and it is only reached when the caller has registered uncorrelated subqueries, which only the physical planner does.So each existing entry point keeps its signature and delegates to a new
_with_subquery_contexttwin that passesSubqueryContext::default(). External callers behave exactly as before (a scalar subquery lowers tonot_impl_err, matching today's behavior when the map was empty); only the physical planner uses the new variants. The trade-off is one extra public function per entry point in exchange for zero breakage of the existing signatures.Are these changes tested?
Existing test coverage (including the
subquerysqllogictests) exercises the changed path. No behavior change is intended.Are there any user-facing changes?
Breaking: the
subquery_indexesandsubquery_resultspublic fields are removed fromdatafusion_expr::execution_props::ExecutionProps. These fields shipped in54.0.0(added in #21240), so this is a breaking change and is documented in the 55.0.0 upgrade guide. Only code that read or wrote those fields directly is affected; in practice only the physical planner populated them. Callers migrate by building aSubqueryContextand using the_with_subquery_contextlowering entry points (see the upgrade guide for a before/after example).Additive, non-breaking: for each existing public lowering entry point in
datafusion-physical-expranddatafusion::physical_planner(create_physical_expr,create_physical_exprs,create_physical_sort_expr[s],create_physical_partitioning,create_window_expr[_with_name]) there is now a_with_subquery_contextsibling.LoweredAggregateBuildergains awith_subquery_contextmethod, andSubqueryContextis a new public type. Every original function keeps its signature and delegates withSubqueryContext::default().