add tests

Author: rluvaton

datafusion/physical-expr-adapter/src/schema_rewriter.rs

@@ -1734,4 +1734,131 @@ mod tests {
         assert_eq!(cast_expr.input_field().data_type(), &DataType::Int32);
         assert_eq!(cast_expr.target_field().data_type(), &DataType::Int64);
     }
+
+    /// A mock expression with an in-scope child and an out-of-scope child.
+    /// Used to verify that scoped traversal does not modify out-of-scope children.
+    #[derive(Debug, Hash, Clone)]
+    struct ScopedExprMock {
+        in_scope_child: Arc<dyn PhysicalExpr>,
+        out_of_scope_child: Arc<dyn PhysicalExpr>,
+    }
+
+    impl PartialEq for ScopedExprMock {
+        fn eq(&self, other: &Self) -> bool {
+            self.in_scope_child.as_ref() == other.in_scope_child.as_ref()
+                && self.out_of_scope_child.as_ref()
+                    == other.out_of_scope_child.as_ref()
+        }
+    }
+
+    impl Eq for ScopedExprMock {}
+
+    impl std::fmt::Display for ScopedExprMock {
+        fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+            write!(
+                f,
+                "scoped_mock({}, {})",
+                self.in_scope_child, self.out_of_scope_child
+            )
+        }
+    }
+
+    impl PhysicalExpr for ScopedExprMock {
+        fn as_any(&self) -> &dyn std::any::Any {
+            self
+        }
+
+        fn return_field(
+            &self,
+            input_schema: &Schema,
+        ) -> Result<Arc<Field>> {
+            self.in_scope_child.return_field(input_schema)
+        }
+
+        fn evaluate(
+            &self,
+            _batch: &RecordBatch,
+        ) -> Result<datafusion_expr::ColumnarValue> {
+            unimplemented!("ScopedExprMock does not support evaluation")
+        }
+
+        fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+            vec![&self.in_scope_child, &self.out_of_scope_child]
+        }
+
+        fn with_new_children(
+            self: Arc<Self>,
+            children: Vec<Arc<dyn PhysicalExpr>>,
+        ) -> Result<Arc<dyn PhysicalExpr>> {
+            assert_eq!(children.len(), 2);
+            let mut iter = children.into_iter();
+            Ok(Arc::new(Self {
+                in_scope_child: iter.next().unwrap(),
+                out_of_scope_child: iter.next().unwrap(),
+            }))
+        }
+
+        fn children_in_scope(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+            vec![&self.in_scope_child]
+        }
+
+        fn with_new_children_in_scope(
+            self: Arc<Self>,
+            children_in_scope: Vec<Arc<dyn PhysicalExpr>>,
+        ) -> Result<Arc<dyn PhysicalExpr>> {
+            assert_eq!(children_in_scope.len(), 1);
+            Ok(Arc::new(Self {
+                in_scope_child: children_in_scope.into_iter().next().unwrap(),
+                out_of_scope_child: Arc::clone(&self.out_of_scope_child),
+            }))
+        }
+
+        fn fmt_sql(
+            &self,
+            f: &mut std::fmt::Formatter<'_>,
+        ) -> std::fmt::Result {
+            std::fmt::Display::fmt(&self, f)
+        }
+    }
+
+    #[test]
+    fn test_replace_columns_with_literals_does_not_modify_out_of_scope_children() {
+        // The in-scope child references column "a" which should be replaced
+        let in_scope_child: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
+        // The out-of-scope child also references a column "a" but should NOT be replaced
+        let out_of_scope_child: Arc<dyn PhysicalExpr> =
+            Arc::new(Column::new("a", 0));
+
+        let expr: Arc<dyn PhysicalExpr> = Arc::new(ScopedExprMock {
+            in_scope_child,
+            out_of_scope_child,
+        });
+
+        let mut replacements = HashMap::new();
+        replacements.insert("a", ScalarValue::Int32(Some(42)));
+
+        let result = replace_columns_with_literals(expr, &replacements).unwrap();
+
+        let mock = result
+            .as_any()
+            .downcast_ref::<ScopedExprMock>()
+            .expect("Should still be ScopedExprMock");
+
+        // The in-scope child "a" should be replaced with literal 42
+        let in_scope_lit = mock
+            .in_scope_child
+            .as_any()
+            .downcast_ref::<Literal>()
+            .expect("in_scope_child should be replaced with Literal");
+        assert_eq!(in_scope_lit.value(), &ScalarValue::Int32(Some(42)));
+
+        // The out-of-scope child "a@0" should be UNCHANGED (still a Column)
+        let out_of_scope_col = mock
+            .out_of_scope_child
+            .as_any()
+            .downcast_ref::<Column>()
+            .expect("out_of_scope_child should still be Column");
+        assert_eq!(out_of_scope_col.name(), "a");
+        assert_eq!(out_of_scope_col.index(), 0);
+    }
 }

datafusion/physical-expr/src/projection.rs

@@ -1200,7 +1200,7 @@ pub(crate) mod tests {
     use super::*;
     use crate::equivalence::{EquivalenceProperties, convert_to_orderings};
     use crate::expressions::{BinaryExpr, col};
-    use crate::utils::tests::TestScalarUDF;
+    use crate::utils::tests::{ScopedExprMock, TestScalarUDF};
     use crate::{PhysicalExprRef, ScalarFunctionExpr};
 
     use arrow::compute::SortOptions;
@@ -3038,4 +3038,162 @@ pub(crate) mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_update_expr_does_not_modify_out_of_scope_children() -> Result<()> {
+        // Outer schema: [a, b, c]
+        // Expression: ScopedExprMock(in_scope=a@0, out_of_scope=x@0)
+        // Projection: [c@2 as c_new, a@0 as a_new, b@1 as b_new]
+        // After unproject: in_scope should become c@2, out_of_scope should stay x@0
+        let in_scope_child: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a_new", 1));
+        let out_of_scope_child: Arc<dyn PhysicalExpr> =
+            Arc::new(Column::new("x", 0));
+
+        let expr: Arc<dyn PhysicalExpr> = Arc::new(ScopedExprMock {
+            in_scope_child,
+            out_of_scope_child,
+        });
+
+        let projected_exprs = vec![
+            ProjectionExpr {
+                expr: Arc::new(Column::new("c", 2)),
+                alias: "c_new".to_string(),
+            },
+            ProjectionExpr {
+                expr: Arc::new(Column::new("a", 0)),
+                alias: "a_new".to_string(),
+            },
+            ProjectionExpr {
+                expr: Arc::new(Column::new("b", 1)),
+                alias: "b_new".to_string(),
+            },
+        ];
+
+        let result =
+            update_expr(&expr, &projected_exprs, true)?.expect("Should be valid");
+
+        let mock = result
+            .as_any()
+            .downcast_ref::<ScopedExprMock>()
+            .expect("Should still be ScopedExprMock");
+
+        // The in-scope child "a_new@1" should be unprojected to "a@0"
+        let in_scope_col = mock
+            .in_scope_child
+            .as_any()
+            .downcast_ref::<Column>()
+            .expect("in_scope_child should be Column");
+        assert_eq!(in_scope_col.name(), "a");
+        assert_eq!(in_scope_col.index(), 0);
+
+        // The out-of-scope child "x@0" should be UNCHANGED
+        let out_of_scope_col = mock
+            .out_of_scope_child
+            .as_any()
+            .downcast_ref::<Column>()
+            .expect("out_of_scope_child should be Column");
+        assert_eq!(out_of_scope_col.name(), "x");
+        assert_eq!(out_of_scope_col.index(), 0);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_project_ordering_does_not_modify_out_of_scope_children() {
+        // Schema: [a: Int32, b: Int32]
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("a", DataType::Int32, false),
+            Field::new("b", DataType::Int32, false),
+        ]));
+
+        let in_scope_child: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
+        let out_of_scope_child: Arc<dyn PhysicalExpr> =
+            Arc::new(Column::new("x", 0));
+
+        let scoped_expr: Arc<dyn PhysicalExpr> = Arc::new(ScopedExprMock {
+            in_scope_child,
+            out_of_scope_child,
+        });
+
+        let ordering =
+            LexOrdering::new(vec![PhysicalSortExpr::new(scoped_expr, SortOptions::new(false, false))])
+                .unwrap();
+
+        let result = project_ordering(&ordering, &schema).expect("Should project");
+
+        let projected_expr = &result.first().expr;
+        let mock = projected_expr
+            .as_any()
+            .downcast_ref::<ScopedExprMock>()
+            .expect("Should still be ScopedExprMock");
+
+        // The in-scope column "a" should be reindexed (stays at 0 in this case)
+        let in_scope_col = mock
+            .in_scope_child
+            .as_any()
+            .downcast_ref::<Column>()
+            .expect("in_scope_child should be Column");
+        assert_eq!(in_scope_col.name(), "a");
+        assert_eq!(in_scope_col.index(), 0);
+
+        // The out-of-scope child "x@0" should be UNCHANGED
+        let out_of_scope_col = mock
+            .out_of_scope_child
+            .as_any()
+            .downcast_ref::<Column>()
+            .expect("out_of_scope_child should be Column");
+        assert_eq!(out_of_scope_col.name(), "x");
+        assert_eq!(out_of_scope_col.index(), 0);
+    }
+
+    #[test]
+    fn test_projection_mapping_does_not_modify_out_of_scope_children() -> Result<()> {
+        // Input schema: [a: Int32, b: Int32]
+        let input_schema = Arc::new(Schema::new(vec![
+            Field::new("a", DataType::Int32, false),
+            Field::new("b", DataType::Int32, false),
+        ]));
+
+        let in_scope_child: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
+        let out_of_scope_child: Arc<dyn PhysicalExpr> =
+            Arc::new(Column::new("x", 0));
+
+        let scoped_expr: Arc<dyn PhysicalExpr> = Arc::new(ScopedExprMock {
+            in_scope_child,
+            out_of_scope_child,
+        });
+
+        // Project: [ScopedExprMock as "result"]
+        let projection_exprs = vec![(scoped_expr, "result".to_string())];
+
+        let mapping = ProjectionMapping::try_new(projection_exprs, &input_schema)?;
+
+        // The source expression in the mapping should have its in-scope column
+        // validated but the out-of-scope column left untouched
+        let (source_expr, _targets) = mapping.iter().next().unwrap();
+        let mock = source_expr
+            .as_any()
+            .downcast_ref::<ScopedExprMock>()
+            .expect("Should still be ScopedExprMock");
+
+        // In-scope child: "a@0" should still be "a@0" (name matches schema)
+        let in_scope_col = mock
+            .in_scope_child
+            .as_any()
+            .downcast_ref::<Column>()
+            .expect("in_scope_child should be Column");
+        assert_eq!(in_scope_col.name(), "a");
+        assert_eq!(in_scope_col.index(), 0);
+
+        // Out-of-scope child: "x@0" should be UNCHANGED
+        let out_of_scope_col = mock
+            .out_of_scope_child
+            .as_any()
+            .downcast_ref::<Column>()
+            .expect("out_of_scope_child should be Column");
+        assert_eq!(out_of_scope_col.name(), "x");
+        assert_eq!(out_of_scope_col.index(), 0);
+
+        Ok(())
+    }
 }