Skip to content

refactor: thread SubqueryContext explicitly through physical planning#22340

Open
timsaucer wants to merge 2 commits into
apache:mainfrom
timsaucer:refactor/subquery-context-explicit
Open

refactor: thread SubqueryContext explicitly through physical planning#22340
timsaucer wants to merge 2 commits into
apache:mainfrom
timsaucer:refactor/subquery-context-explicit

Conversation

@timsaucer

@timsaucer timsaucer commented May 18, 2026

Copy link
Copy Markdown
Member

Which issue does this PR close?

None, but it is a precursor to making QueryPlanner/PhysicalPlanner take &dyn Session instead 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 Expr into a physical Arc<dyn PhysicalExpr> — the job of create_physical_expr in physical-expr/src/planner.rs. That function only receives &ExecutionProps, so any state it needs has to arrive through ExecutionProps.

Uncorrelated scalar subqueries need such state: the Expr::ScalarSubquery branch of create_physical_expr has to look up which result slot a given subquery maps to. Today DefaultPhysicalPlanner::create_initial_plan supplies that by cloning the SessionState and mutating the embedded ExecutionProps:

let mut owned = session_state.clone();
owned.execution_props_mut().subquery_indexes = index_map;
owned.execution_props_mut().subquery_results = results.clone();

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:

  • Session is a trait object with no clone(), so session_state.clone() has no equivalent — you cannot cheaply produce an owned copy to mutate.
  • Session::execution_props(&self) returns &ExecutionProps immutably by design (it is called from &self TableProvider hot paths), so there is no execution_props_mut() to write into.

So the subquery state cannot ride inside ExecutionProps under &dyn Session. It has to be threaded explicitly instead. This PR does that threading now, decoupled from the larger &dyn Session change.

What changes are included in this PR?

  • Adds SubqueryContext in datafusion-expr — a small carrier bundling the Subquery -> SubqueryIndex map and the shared ScalarSubqueryResults.
  • Threads &SubqueryContext explicitly from DefaultPhysicalPlanner down into expression lowering, replacing the ExecutionProps side channel.
  • Removes the SessionState::clone() + execution_props_mut() mutation at the write site.
  • Removes the subquery_indexes and subquery_results fields from ExecutionProps (breaking — see user-facing changes).

Why add _with_subquery_context variants instead of changing the existing signatures?

Adding a 4th subquery_ctx parameter to create_physical_expr is a cleaner end state but has the wrong cost:

  • create_physical_expr has ~92 callers in this repo plus downstream crates (datafusion-comet, ballista, datafusion-python, custom planners). All would have to pass the new argument.
  • None of them need scalar-subquery support. Only one read site cares — the Expr::ScalarSubquery branch — 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_context twin that passes SubqueryContext::default(). External callers behave exactly as before (a scalar subquery lowers to not_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 subquery sqllogictests) exercises the changed path. No behavior change is intended.

Are there any user-facing changes?

Breaking: the subquery_indexes and subquery_results public fields are removed from datafusion_expr::execution_props::ExecutionProps. These fields shipped in 54.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 a SubqueryContext and using the _with_subquery_context lowering entry points (see the upgrade guide for a before/after example).

Additive, non-breaking: for each existing public lowering entry point in datafusion-physical-expr and datafusion::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_context sibling. LoweredAggregateBuilder gains a with_subquery_context method, and SubqueryContext is a new public type. Every original function keeps its signature and delegates with SubqueryContext::default().

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>
@github-actions github-actions Bot added logical-expr Logical plan and expressions physical-expr Changes to the physical-expr crates core Core DataFusion crate labels May 18, 2026
@timsaucer

Copy link
Copy Markdown
Member Author

@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 &dyn Session instead of &SessionState in the QueryPlanner.

@github-actions

github-actions Bot commented May 18, 2026

Copy link
Copy Markdown

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
     Cloning apache/main
    Building datafusion v54.0.0 (current)
       Built [ 108.942s] (current)
     Parsing datafusion v54.0.0 (current)
      Parsed [   0.038s] (current)
    Building datafusion v54.0.0 (baseline)
       Built [ 105.903s] (baseline)
     Parsing datafusion v54.0.0 (baseline)
      Parsed [   0.039s] (baseline)
    Checking datafusion v54.0.0 -> v54.0.0 (no change; assume patch)
     Checked [   0.899s] 223 checks: 223 pass, 30 skip
     Summary no semver update required
    Finished [ 217.672s] datafusion
    Building datafusion-expr v54.0.0 (current)
       Built [  26.544s] (current)
     Parsing datafusion-expr v54.0.0 (current)
      Parsed [   0.080s] (current)
    Building datafusion-expr v54.0.0 (baseline)
       Built [  26.523s] (baseline)
     Parsing datafusion-expr v54.0.0 (baseline)
      Parsed [   0.084s] (baseline)
    Checking datafusion-expr v54.0.0 -> v54.0.0 (no change; assume patch)
     Checked [   2.099s] 223 checks: 222 pass, 1 fail, 0 warn, 30 skip

--- failure struct_pub_field_missing: pub struct's pub field removed or renamed ---

Description:
A publicly-visible struct has at least one public field that is no longer available under its prior name. It may have been renamed or removed entirely.
        ref: https://doc.rust-lang.org/cargo/reference/semver.html#item-remove
       impl: https://github.com/obi1kenobi/cargo-semver-checks/tree/v0.48.0/src/lints/struct_pub_field_missing.ron

Failed in:
  field subquery_indexes of struct ExecutionProps, previously in file /home/runner/work/datafusion/datafusion/target/semver-checks/git-apache_main/257152542e7c03dcd1162d720af457922e203cca/datafusion/expr/src/execution_props.rs:69
  field subquery_results of struct ExecutionProps, previously in file /home/runner/work/datafusion/datafusion/target/semver-checks/git-apache_main/257152542e7c03dcd1162d720af457922e203cca/datafusion/expr/src/execution_props.rs:72

     Summary semver requires new major version: 1 major and 0 minor checks failed
    Finished [  56.324s] datafusion-expr
    Building datafusion-physical-expr v54.0.0 (current)
       Built [  28.200s] (current)
     Parsing datafusion-physical-expr v54.0.0 (current)
      Parsed [   0.048s] (current)
    Building datafusion-physical-expr v54.0.0 (baseline)
       Built [  28.061s] (baseline)
     Parsing datafusion-physical-expr v54.0.0 (baseline)
      Parsed [   0.049s] (baseline)
    Checking datafusion-physical-expr v54.0.0 -> v54.0.0 (no change; assume patch)
     Checked [   0.517s] 223 checks: 223 pass, 30 skip
     Summary no semver update required
    Finished [  58.143s] datafusion-physical-expr

@github-actions github-actions Bot added the auto detected api change Auto detected API change label May 18, 2026
@timsaucer timsaucer marked this pull request as ready for review June 17, 2026 15:20
@timsaucer

Copy link
Copy Markdown
Member Author

@milenkovicm this is the PR I was talking about to unblock FFI query planner

@milenkovicm

Copy link
Copy Markdown
Contributor

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
@timsaucer timsaucer requested a review from gabotechs July 8, 2026 16:59
Comment on lines +247 to 253
pub fn create_physical_sort_exprs_with_subquery_context(
exprs: &[SortExpr],
input_dfschema: &DFSchema,
execution_props: &ExecutionProps,
subquery_ctx: &SubqueryContext,
) -> Result<Vec<PhysicalSortExpr>> {
exprs

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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...

Comment on lines -475 to +480
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());

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I gave it a try in a draft:

And it looks dead simple, but not sure if I'm missing something.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

auto detected api change Auto detected API change core Core DataFusion crate logical-expr Logical plan and expressions physical-expr Changes to the physical-expr crates

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants