add tree scope

rluvaton · rluvaton · commit 93807d4c9fa8 · 2026-04-05T16:01:07.000+03:00
diff --git a/datafusion/physical-expr-adapter/src/schema_rewriter.rs b/datafusion/physical-expr-adapter/src/schema_rewriter.rs
@@ -26,6 +26,7 @@ use std::sync::Arc;
 
 use arrow::array::RecordBatch;
 use arrow::datatypes::{DataType, Field, FieldRef, SchemaRef};
+use datafusion_common::tree_node::ScopedTreeNode;
 use datafusion_common::{
     DataFusionError, Result, ScalarValue, exec_err,
     metadata::FieldMetadata,
@@ -69,7 +70,7 @@ where
     K: Borrow<str> + Eq + Hash,
     V: Borrow<ScalarValue>,
 {
-    expr.transform_down(|expr| {
+    expr.transform_down_in_scope(|expr| {
         if let Some(column) = expr.as_any().downcast_ref::<Column>()
             && let Some(replacement_value) = replacements.get(column.name())
         {
diff --git a/datafusion/physical-expr-common/src/physical_expr.rs b/datafusion/physical-expr-common/src/physical_expr.rs
@@ -165,12 +165,37 @@ pub trait PhysicalExpr: Any + Send + Sync + Display + Debug + DynEq + DynHash {
     /// Get a list of child PhysicalExpr that provide the input for this expr.
     fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>>;
 
+    /// Get a list of child PhysicalExpr that provide the input for this expr that are in the same scope as this expression.
+    /// 
+    /// Due to the majority of expressions being in the same scope the default implementation is to call to [`Self::children`] 
+    /// 
+    /// To know if specific child is considered in the same scope you can answer this simple question:
+    /// If that child is a `Column` would that column can be evaluated with the same input schema
+    /// Expressions like `plus`, `sum`, etc have all children in scope.
+    /// Lambda expressions like `array_filter(list, value -> value + 1)`, have the `list` in the same scope and the lambda function in different scope
+    fn children_in_scope(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+        self.children()
+    }
+
     /// Returns a new PhysicalExpr where all children were replaced by new exprs.
     fn with_new_children(
         self: Arc<Self>,
         children: Vec<Arc<dyn PhysicalExpr>>,
     ) -> Result<Arc<dyn PhysicalExpr>>;
 
+    /// Returns a new PhysicalExpr where all scoped children were replaced by new exprs.
+    /// 
+    /// See [`Self::children_in_scope`] for definition of what child considered a scope
+    /// 
+    /// Due to the majority of expressions being in the same scope the default implementation is to call to [`Self::with_new_children`] 
+    /// 
+    fn with_new_children_in_scope(
+        self: Arc<Self>,
+        children_in_scope: Vec<Arc<dyn PhysicalExpr>>,
+    ) -> Result<Arc<dyn PhysicalExpr>> {
+        self.with_new_children(children_in_scope)
+    }
+
     /// Computes the output interval for the expression, given the input
     /// intervals.
     ///
@@ -476,16 +501,40 @@ pub fn with_new_children_if_necessary(
     );
 
     if children.is_empty()
-        || children
-            .iter()
-            .zip(old_children.iter())
-            .any(|(c1, c2)| !Arc::ptr_eq(c1, c2))
+      || children
+      .iter()
+      .zip(old_children.iter())
+      .any(|(c1, c2)| !Arc::ptr_eq(c1, c2))
     {
         Ok(expr.with_new_children(children)?)
     } else {
         Ok(expr)
     }
 }
+/// Returns a copy of this expr if we change any child according to the pointer comparison.
+/// The size of `children_in_scope` must be equal to the size of [`PhysicalExpr::children_in_scope()`].
+pub fn with_new_children_in_scope_if_necessary(
+    expr: Arc<dyn PhysicalExpr>,
+    children_in_scope: Vec<Arc<dyn PhysicalExpr>>,
+) -> Result<Arc<dyn PhysicalExpr>> {
+    let old_children_in_scope = expr.children_in_scope();
+    assert_eq_or_internal_err!(
+        children_in_scope.len(),
+        old_children_in_scope.len(),
+        "PhysicalExpr: Wrong number of children in scope"
+    );
+
+    if children_in_scope.is_empty()
+      || children_in_scope
+      .iter()
+      .zip(old_children_in_scope.iter())
+      .any(|(c1, c2)| !Arc::ptr_eq(c1, c2))
+    {
+        Ok(expr.with_new_children_in_scope(children_in_scope)?)
+    } else {
+        Ok(expr)
+    }
+}
 
 /// Returns [`Display`] able a list of [`PhysicalExpr`]
 ///
diff --git a/datafusion/physical-expr-common/src/tree_node.rs b/datafusion/physical-expr-common/src/tree_node.rs
@@ -20,10 +20,10 @@
 use std::fmt::{self, Display, Formatter};
 use std::sync::Arc;
 
-use crate::physical_expr::{PhysicalExpr, with_new_children_if_necessary};
+use crate::physical_expr::{PhysicalExpr, with_new_children_if_necessary, with_new_children_in_scope_if_necessary};
 
 use datafusion_common::Result;
-use datafusion_common::tree_node::{ConcreteTreeNode, DynTreeNode};
+use datafusion_common::tree_node::{ConcreteTreeNode, DynScopedTreeNode, DynTreeNode};
 
 impl DynTreeNode for dyn PhysicalExpr {
     fn arc_children(&self) -> Vec<&Arc<Self>> {
@@ -39,6 +39,20 @@ impl DynTreeNode for dyn PhysicalExpr {
     }
 }
 
+impl DynScopedTreeNode for dyn PhysicalExpr {
+    fn arc_children_in_scope(&self) -> Vec<&Arc<Self>> {
+        self.children_in_scope()
+    }
+
+    fn with_new_arc_children_in_scope(
+        &self,
+        arc_self: Arc<Self>,
+        new_children: Vec<Arc<Self>>,
+    ) -> Result<Arc<Self>> {
+        with_new_children_in_scope_if_necessary(arc_self, new_children)
+    }
+}
+
 /// A node object encapsulating a [`PhysicalExpr`] node with a payload. Since there are
 /// two ways to access child plans—directly from the plan  and through child nodes—it's
 /// recommended to perform mutable operations via [`Self::update_expr_from_children`].
diff --git a/datafusion/physical-expr/src/expressions/case.rs b/datafusion/physical-expr/src/expressions/case.rs
@@ -40,7 +40,9 @@ use std::{any::Any, sync::Arc};
 
 use crate::expressions::case::literal_lookup_table::LiteralLookupTable;
 use arrow::compute::kernels::merge::{MergeIndex, merge, merge_n};
-use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
+use datafusion_common::tree_node::{
+    ScopedTreeNode, Transformed, TreeNode, TreeNodeRecursion,
+};
 use datafusion_physical_expr_common::datum::compare_with_eq;
 use datafusion_physical_expr_common::utils::scatter;
 use itertools::Itertools;
@@ -130,7 +132,7 @@ impl CaseBody {
         // Determine the set of columns that are used in all the expressions of the case body.
         let mut used_column_indices = IndexSet::<usize>::new();
         let mut collect_column_indices = |expr: &Arc<dyn PhysicalExpr>| {
-            expr.apply(|expr| {
+            expr.apply_in_scope(|expr| {
                 if let Some(column) = expr.as_any().downcast_ref::<Column>() {
                     used_column_indices.insert(column.index());
                 }
@@ -161,7 +163,7 @@ impl CaseBody {
         // using the column index mapping.
         let project = |expr: &Arc<dyn PhysicalExpr>| -> Result<Arc<dyn PhysicalExpr>> {
             Arc::clone(expr)
-                .transform_down(|e| {
+                .transform_down_in_scope(|e| {
                     if let Some(column) = e.as_any().downcast_ref::<Column>() {
                         let original = column.index();
                         let projected = *column_index_map.get(&original).unwrap();
@@ -1397,7 +1399,7 @@ fn replace_with_null(
     input_schema: &Schema,
 ) -> Result<Arc<dyn PhysicalExpr>, DataFusionError> {
     let with_null = Arc::clone(expr)
-        .transform_down(|e| {
+        .transform_down_in_scope(|e| {
             if e.as_ref().dyn_eq(expr_to_replace) {
                 let data_type = e.data_type(input_schema)?;
                 let null_literal = lit(ScalarValue::try_new_null(&data_type)?);
@@ -1928,135 +1930,6 @@ mod tests {
         Ok(())
     }
 
-    #[test]
-    fn case_without_expr_and_with_custom_column_impl() -> Result<()> {
-        /// Represents the column at a given index in a RecordBatch that is inside a Spark lambda function
-        ///
-        /// This is the same as the datafusion [`datafusion::physical_expr::expressions::Column`] except that it store the entire info so that it can be used in lambda execution
-        #[derive(Debug, Hash, PartialEq, Eq, Clone)]
-        pub struct CustomColumn {
-            /// The name of the column (used for debugging and display purposes)
-            name: String,
-            /// The index of the column in its schema
-            index: usize,
-            data_type: DataType,
-            nullable: bool,
-        }
-
-        impl CustomColumn {
-            pub fn new_with_schema(
-                name: &str,
-                schema: &Schema,
-            ) -> Result<Arc<dyn PhysicalExpr>> {
-                let index = schema.index_of(name)?;
-                let field = schema.field(index);
-                Ok(Arc::new(CustomColumn {
-                    name: name.to_string(),
-                    index,
-                    data_type: field.data_type().clone(),
-                    nullable: field.is_nullable(),
-                }))
-            }
-        }
-
-        impl std::fmt::Display for CustomColumn {
-            fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
-                write!(f, "{}@{}", self.name, self.index)
-            }
-        }
-
-        impl PhysicalExpr for CustomColumn {
-            fn as_any(&self) -> &dyn Any {
-                self
-            }
-
-            fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
-                Ok(self.data_type.clone())
-            }
-
-            fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
-                Ok(self.nullable)
-            }
-
-            fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
-                self.bounds_check(batch.schema().as_ref())?;
-                Ok(ColumnarValue::Array(Arc::clone(batch.column(self.index))))
-            }
-
-            fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
-                vec![]
-            }
-
-            fn with_new_children(
-                self: Arc<Self>,
-                _children: Vec<Arc<dyn PhysicalExpr>>,
-            ) -> Result<Arc<dyn PhysicalExpr>> {
-                Ok(self)
-            }
-
-            fn fmt_sql(&self, _: &mut Formatter<'_>) -> std::fmt::Result {
-                unimplemented!()
-            }
-        }
-
-        impl CustomColumn {
-            fn bounds_check(&self, input_schema: &Schema) -> Result<()> {
-                if self.index < input_schema.fields.len() {
-                    Ok(())
-                } else {
-                    internal_err!(
-                        "PhysicalExpr BoundLambdaColumn references column '{}' at index {} (zero-based) but input schema only has {} columns: {:?}",
-                        self.name,
-                        self.index,
-                        input_schema.fields.len(),
-                        input_schema
-                            .fields()
-                            .iter()
-                            .map(|f| f.name())
-                            .collect::<Vec<_>>()
-                    )
-                }
-            }
-        }
-
-        let batch = case_test_batch()?;
-        let schema = batch.schema();
-
-        // CASE WHEN a = 'foo' THEN 123 WHEN a = 'bar' THEN 456 END
-        let when1 = binary(
-            CustomColumn::new_with_schema("a", &schema)?,
-            Operator::Eq,
-            lit("foo"),
-            &batch.schema(),
-        )?;
-        let then1 = lit(123i32);
-        let when2 = binary(
-            CustomColumn::new_with_schema("a", &schema)?,
-            Operator::Eq,
-            lit("bar"),
-            &batch.schema(),
-        )?;
-        let then2 = lit(456i32);
-
-        let expr = generate_case_when_with_type_coercion(
-            None,
-            vec![(when1, then1), (when2, then2)],
-            None,
-            schema.as_ref(),
-        )?;
-        let result = expr
-            .evaluate(&batch)?
-            .into_array(batch.num_rows())
-            .expect("Failed to convert to array");
-        let result = as_int32_array(&result)?;
-
-        let expected = &Int32Array::from(vec![Some(123), None, None, Some(456)]);
-
-        assert_eq!(expected, result);
-
-        Ok(())
-    }
-
     #[test]
     fn case_with_expr_when_null() -> Result<()> {
         let batch = case_test_batch()?;
@@ -2552,7 +2425,7 @@ mod tests {
             .unwrap();
 
         let expr3 = Arc::clone(&expr)
-            .transform_down(|e| {
+            .transform_down_in_scope(|e| {
                 let transformed = match e.as_any().downcast_ref::<Literal>() {
                     Some(lit_value) => match lit_value.value() {
                         ScalarValue::Utf8(Some(str_value)) => {
diff --git a/datafusion/physical-expr/src/physical_expr.rs b/datafusion/physical-expr/src/physical_expr.rs
@@ -22,7 +22,7 @@ use crate::{LexOrdering, PhysicalSortExpr, create_physical_expr};
 
 use arrow::compute::SortOptions;
 use arrow::datatypes::{Schema, SchemaRef};
-use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_common::tree_node::{ScopedTreeNode, Transformed, TransformedResult, TreeNode};
 use datafusion_common::{DFSchema, HashMap};
 use datafusion_common::{Result, plan_err};
 use datafusion_expr::execution_props::ExecutionProps;
@@ -38,7 +38,7 @@ pub fn add_offset_to_expr(
     expr: Arc<dyn PhysicalExpr>,
     offset: isize,
 ) -> Result<Arc<dyn PhysicalExpr>> {
-    expr.transform_down(|e| match e.as_any().downcast_ref::<Column>() {
+    expr.transform_down_in_scope(|e| match e.as_any().downcast_ref::<Column>() {
         Some(col) => {
             let Some(idx) = col.index().checked_add_signed(offset) else {
                 return plan_err!("Column index overflow");
diff --git a/datafusion/physical-expr/src/projection.rs b/datafusion/physical-expr/src/projection.rs
@@ -27,7 +27,7 @@ use crate::utils::collect_columns;
 use arrow::array::{RecordBatch, RecordBatchOptions};
 use arrow::datatypes::{Field, Schema, SchemaRef};
 use datafusion_common::stats::{ColumnStatistics, Precision};
-use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_common::tree_node::{ScopedTreeNode, Transformed, TransformedResult, TreeNode};
 use datafusion_common::{
     Result, ScalarValue, Statistics, assert_or_internal_err, internal_datafusion_err,
     plan_err,
@@ -920,7 +920,7 @@ pub fn update_expr(
     let mut state = RewriteState::Unchanged;
 
     let new_expr = Arc::clone(expr)
-        .transform_up(|expr| {
+        .transform_up_in_scope(|expr| {
             if state == RewriteState::RewrittenInvalid {
                 return Ok(Transformed::no(expr));
             }
@@ -1043,7 +1043,7 @@ impl ProjectionMapping {
         let mut map = IndexMap::<_, ProjectionTargets>::new();
         for (expr_idx, (expr, name)) in expr.into_iter().enumerate() {
             let target_expr = Arc::new(Column::new(&name, expr_idx)) as _;
-            let source_expr = expr.transform_down(|e| match e.as_any().downcast_ref::<Column>() {
+            let source_expr = expr.transform_down_in_scope(|e| match e.as_any().downcast_ref::<Column>() {
                 Some(col) => {
                     // Sometimes, an expression and its name in the input_schema
                     // doesn't match. This can cause problems, so we make sure
@@ -1162,7 +1162,7 @@ pub fn project_ordering(
 ) -> Option<LexOrdering> {
     let mut projected_exprs = vec![];
     for PhysicalSortExpr { expr, options } in ordering.iter() {
-        let transformed = Arc::clone(expr).transform_up(|expr| {
+        let transformed = Arc::clone(expr).transform_up_in_scope(|expr| {
             let Some(col) = expr.as_any().downcast_ref::<Column>() else {
                 return Ok(Transformed::no(expr));
             };
diff --git a/datafusion/physical-expr/src/utils/mod.rs b/datafusion/physical-expr/src/utils/mod.rs
@@ -16,6 +16,7 @@
 // under the License.
 
 mod guarantee;
+use datafusion_common::tree_node::ScopedTreeNode;
 pub use guarantee::{Guarantee, LiteralGuarantee};
 
 use std::borrow::Borrow;
@@ -251,7 +252,7 @@ pub fn reassign_expr_columns(
     expr: Arc<dyn PhysicalExpr>,
     schema: &Schema,
 ) -> Result<Arc<dyn PhysicalExpr>> {
-    expr.transform_down(|expr| {
+    expr.transform_down_in_scope(|expr| {
         if let Some(column) = expr.as_any().downcast_ref::<Column>() {
             let index = schema.index_of(column.name())?;
 
diff --git a/datafusion/physical-plan/src/filter_pushdown.rs b/datafusion/physical-plan/src/filter_pushdown.rs
