From 3b38dfa0752c36f5a4131e22ab0ad16fdc2e3e34 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Florian=20M=C3=BCller?= Date: Wed, 29 Jul 2026 21:49:59 +0200 Subject: [PATCH] refactor: move lambda evaluation into Physical Planning Context --- datafusion/expr/src/execution_props.rs | 24 +------ .../expr/src/physical_planning_context.rs | 65 ++++++++++++++++--- datafusion/physical-expr/src/planner.rs | 28 ++++---- .../library-user-guide/upgrading/55.0.0.md | 17 ++++- 4 files changed, 89 insertions(+), 45 deletions(-) diff --git a/datafusion/expr/src/execution_props.rs b/datafusion/expr/src/execution_props.rs index 9910918c6ea2a..7c5369d1144dd 100644 --- a/datafusion/expr/src/execution_props.rs +++ b/datafusion/expr/src/execution_props.rs @@ -18,7 +18,6 @@ use crate::var_provider::{VarProvider, VarType}; use chrono::{DateTime, Utc}; use datafusion_common::HashMap; -use datafusion_common::TableReference; use datafusion_common::alias::AliasGenerator; use datafusion_common::config::ConfigOptions; use std::sync::Arc; @@ -60,10 +59,6 @@ pub struct ExecutionProps { pub config_options: Option>, /// Providers for scalar variables pub var_providers: Option>>, - /// Maps each lambda variable name to its lambda qualifier generated - /// during physical planning. Populated by the physical planner for - /// each lambda before calling `create_physical_expr`. - pub lambda_variable_qualifier: HashMap, } impl Default for ExecutionProps { @@ -80,7 +75,6 @@ impl ExecutionProps { alias_generator: Arc::new(AliasGenerator::new()), config_options: None, var_providers: None, - lambda_variable_qualifier: HashMap::new(), } } @@ -139,22 +133,6 @@ impl ExecutionProps { pub fn config_options(&self) -> Option<&Arc> { self.config_options.as_ref() } - - /// Adds a mapping for each variable to the given qualifier. Existing - /// variables with conflicting names get's shadowed - pub fn with_qualified_lambda_variables( - mut self, - qualifier: &TableReference, - variables: &[String], - ) -> Self { - for var in variables { - self.lambda_variable_qualifier - .entry_ref(var) - .insert(qualifier.clone()); - } - - self - } } #[cfg(test)] @@ -165,7 +143,7 @@ mod test { fn debug() { let props = ExecutionProps::new(); assert_eq!( - "ExecutionProps { query_execution_start_time: None, alias_generator: AliasGenerator { next_id: 1 }, config_options: None, var_providers: None, lambda_variable_qualifier: {} }", + "ExecutionProps { query_execution_start_time: None, alias_generator: AliasGenerator { next_id: 1 }, config_options: None, var_providers: None }", format!("{props:?}") ); } diff --git a/datafusion/expr/src/physical_planning_context.rs b/datafusion/expr/src/physical_planning_context.rs index b1ba63e0718f5..fca6232eb35cd 100644 --- a/datafusion/expr/src/physical_planning_context.rs +++ b/datafusion/expr/src/physical_planning_context.rs @@ -19,32 +19,42 @@ use std::fmt; use std::hash::{Hash, Hasher}; use std::sync::{Arc, Mutex}; -use datafusion_common::{HashMap, Result, ScalarValue, internal_err}; +use datafusion_common::{HashMap, Result, ScalarValue, TableReference, internal_err}; /// Context used while converting a logical plan subtree into a physical plan. /// /// Unlike [`ExecutionProps`](crate::execution_props::ExecutionProps), which /// applies to the overall planning and execution of a query, this context can -/// differ between recursively planned subtrees. It currently carries the state -/// needed to create physical expressions for [`Expr::ScalarSubquery`] nodes -/// that read from a shared -/// [`ScalarSubqueryResults`] container. +/// differ between recursively planned subtrees. It currently carries: +/// +/// * the state needed to create physical expressions for +/// [`Expr::ScalarSubquery`] nodes that read from a shared +/// [`ScalarSubqueryResults`] container, and +/// * the qualifiers assigned to the [`Expr::LambdaVariable`]s that are in scope. /// /// The physical planner builds this context from the set of uncorrelated scalar /// subqueries it has scheduled for a subtree. It is then passed explicitly /// through `create_physical_expr` so that function can find the slot index for -/// each [`Subquery`]. +/// each [`Subquery`]. While planning the body of a lambda, +/// `create_physical_expr` extends the context with the lambda's parameters via +/// [`Self::with_qualified_lambda_variables`]. /// /// An empty [`PhysicalPlanningContext`] (the [`Default`]) is what every /// non-physical-planner caller passes; if such a caller encounters a scalar /// subquery, `create_physical_expr` returns a `not_impl_err`. /// /// [`Expr::ScalarSubquery`]: crate::Expr::ScalarSubquery +/// [`Expr::LambdaVariable`]: crate::Expr::LambdaVariable /// [`Subquery`]: crate::logical_plan::Subquery #[derive(Clone, Debug, Default)] pub struct PhysicalPlanningContext { - indexes: HashMap, + /// Behind an `Arc` because the context is cloned for each lambda body that + /// is planned, and the indexes are the same for the whole subtree. + indexes: Arc>, results: ScalarSubqueryResults, + /// Maps each lambda variable name in scope to the qualifier generated for + /// its lambda during physical planning. + lambda_variable_qualifier: HashMap, } impl PhysicalPlanningContext { @@ -55,7 +65,11 @@ impl PhysicalPlanningContext { indexes: HashMap, results: ScalarSubqueryResults, ) -> Self { - Self { indexes, results } + Self { + indexes: Arc::new(indexes), + results, + lambda_variable_qualifier: HashMap::new(), + } } /// Returns the slot index assigned to `subquery`, if any. @@ -70,6 +84,27 @@ impl PhysicalPlanningContext { pub fn results(&self) -> &ScalarSubqueryResults { &self.results } + + /// Adds a mapping for each variable to the given qualifier. Existing + /// variables with conflicting names get shadowed + pub fn with_qualified_lambda_variables( + mut self, + qualifier: &TableReference, + variables: &[String], + ) -> Self { + for var in variables { + self.lambda_variable_qualifier + .entry_ref(var) + .insert(qualifier.clone()); + } + + self + } + + /// Returns the qualifier of the lambda variable `name`, if it is in scope. + pub fn lambda_variable_qualifier(&self, name: &str) -> Option<&TableReference> { + self.lambda_variable_qualifier.get(name) + } } /// Index of a scalar subquery within a [`ScalarSubqueryResults`] container. @@ -192,6 +227,20 @@ mod tests { Ok(()) } + #[test] + fn lambda_variables_shadow_outer_scope() { + let outer = TableReference::bare("lambda_1"); + let inner = TableReference::bare("lambda_2"); + + let ctx = PhysicalPlanningContext::default() + .with_qualified_lambda_variables(&outer, &["x".to_string(), "y".to_string()]) + .with_qualified_lambda_variables(&inner, &["y".to_string()]); + + assert_eq!(ctx.lambda_variable_qualifier("x"), Some(&outer)); + assert_eq!(ctx.lambda_variable_qualifier("y"), Some(&inner)); + assert_eq!(ctx.lambda_variable_qualifier("z"), None); + } + #[test] fn scalar_subquery_results_clear() -> Result<()> { let results = ScalarSubqueryResults::new(1); diff --git a/datafusion/physical-expr/src/planner.rs b/datafusion/physical-expr/src/planner.rs index 3cdd64f7a70d8..f80d1b15bdc59 100644 --- a/datafusion/physical-expr/src/planner.rs +++ b/datafusion/physical-expr/src/planner.rs @@ -121,9 +121,11 @@ use datafusion_expr::{ /// to qualified or unqualified fields by name. /// * `execution_props` - Per-execution properties such as the query start time. /// * `planning_ctx` - The [`PhysicalPlanningContext`] used to resolve -/// `Expr::ScalarSubquery` nodes. The physical planner threads the subquery -/// index map and shared results container from its `ScalarSubqueryExec` -/// construction into calls to `create_physical_expr`. Callers creating +/// `Expr::ScalarSubquery` and `Expr::LambdaVariable` nodes. The physical +/// planner threads the subquery index map and shared results container from +/// its `ScalarSubqueryExec` construction into calls to +/// `create_physical_expr`; the lambda variable qualifiers are added by this +/// function itself as it descends into lambda bodies. Callers creating /// physical expressions outside of physical planning should pass /// `&PhysicalPlanningContext::default()`; converting a scalar subquery then returns a /// planning error. @@ -612,15 +614,15 @@ pub fn create_physical_expr( input_dfschema.metadata().clone(), )?; - let execution_props = execution_props + let planning_ctx = planning_ctx .clone() .with_qualified_lambda_variables(&qualifier, &lambda.params); create_physical_expr( arg, &lambda_schema, - &execution_props, - planning_ctx, + execution_props, + &planning_ctx, ) } _ => create_physical_expr( @@ -657,12 +659,14 @@ pub fn create_physical_expr( plan_datafusion_err!("unresolved LambdaVariable {name}") })?; - let qualifier = execution_props - .lambda_variable_qualifier - .get(name) - .ok_or_else(|| { - plan_datafusion_err!("qualifier for lambda variable {name} not found") - })?; + let qualifier = + planning_ctx + .lambda_variable_qualifier(name) + .ok_or_else(|| { + plan_datafusion_err!( + "qualifier for lambda variable {name} not found" + ) + })?; let index = input_dfschema .index_of_column_by_name(Some(qualifier), name) diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md index 6097c8dc717df..b7ea37523b2ef 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -724,7 +724,7 @@ unit is truncated -- so `time(s) + interval '1 nanosecond'` is a no-op. See [PR #23279](https://github.com/apache/datafusion/pull/23279) for details. -### Scalar-subquery state moved to an explicit `PhysicalPlanningContext` +### Physical-planning state moved to an explicit `PhysicalPlanningContext` The `subquery_indexes` and `subquery_results` public fields on `datafusion_expr::execution_props::ExecutionProps` have been removed. They were @@ -732,6 +732,11 @@ added in `54.0.0` as the channel through which the physical planner passed uncorrelated scalar-subquery state to functions that create physical `Arc` values from logical `Expr` values. +The `lambda_variable_qualifier` public field and the +`with_qualified_lambda_variables` method on `ExecutionProps` have been removed +for the same reason: they carried the qualifiers of the lambda variables in +scope while `create_physical_expr` descended into a lambda body. + That state is now carried by a dedicated `datafusion_expr::physical_planning_context::PhysicalPlanningContext` passed explicitly through functions and planner traits. Unlike `ExecutionProps`, which applies @@ -774,6 +779,13 @@ Convenience methods such as `SessionContext::create_physical_expr` and parameter and forward it. - Code that read or wrote `execution_props.subquery_indexes` / `execution_props.subquery_results`: build a `PhysicalPlanningContext` instead. +- Code that read `execution_props.lambda_variable_qualifier` or called + `ExecutionProps::with_qualified_lambda_variables`: nothing to migrate. + `create_physical_expr` populates the lambda qualifiers itself as it descends + into lambda bodies, so callers planning a `HigherOrderFunction` do not need to + do anything. The equivalent state now lives on `PhysicalPlanningContext`, + reachable via `PhysicalPlanningContext::lambda_variable_qualifier` and + `PhysicalPlanningContext::with_qualified_lambda_variables`. **Migration guide:** @@ -817,7 +829,8 @@ async fn plan_extension( } ``` -See [PR #23649](https://github.com/apache/datafusion/pull/23649) for details. +See [PR #23649](https://github.com/apache/datafusion/pull/23649) and +[PR #23989](https://github.com/apache/datafusion/pull/23989) for details. ### Catalog traits moved to `datafusion-session`