refactor: move lambda evaluation into Physical Planning Context

Author: sweb

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<Arc<ConfigOptions>>,
     /// Providers for scalar variables
     pub var_providers: Option<HashMap<VarType, Arc<dyn VarProvider + Send + Sync>>>,
-    /// 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<String, TableReference>,
 }
 
 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<ConfigOptions>> {
         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:?}")
         );
     }

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<crate::logical_plan::Subquery, SubqueryIndex>,
+    // 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<HashMap<crate::logical_plan::Subquery, SubqueryIndex>>,
     results: ScalarSubqueryResults,
+    /// Maps each lambda variable name in scope to the qualifier generated for
+    /// its lambda during physical planning.
+    lambda_variable_qualifier: HashMap<String, TableReference>,
 }
 
 impl PhysicalPlanningContext {
@@ -55,7 +65,11 @@ impl PhysicalPlanningContext {
         indexes: HashMap<crate::logical_plan::Subquery, SubqueryIndex>,
         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);

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)

docs/source/library-user-guide/upgrading/55.0.0.md

@@ -724,14 +724,19 @@ 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
 added in `54.0.0` as the channel through which the physical planner passed
 uncorrelated scalar-subquery state to functions that create physical
 `Arc<dyn PhysicalExpr>` 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,12 @@ 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`: use
+  `PhysicalPlanningContext::lambda_variable_qualifier` and
+  `PhysicalPlanningContext::with_qualified_lambda_variables`. Note that
+  `create_physical_expr` populates the lambda qualifiers itself, so callers
+  planning a `HigherOrderFunction` do not need to do anything.
 
 **Migration guide:**
 
@@ -817,7 +828,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
+[issue #23700](https://github.com/apache/datafusion/issues/23700) for details.
 
 ### Catalog traits moved to `datafusion-session`