datafusion_physical_plan/
column_rewriter.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::sync::Arc;
19
20use datafusion_common::{
21    DataFusionError, HashMap,
22    tree_node::{Transformed, TreeNodeRecursion, TreeNodeRewriter},
23};
24use datafusion_physical_expr::{PhysicalExpr, expressions::Column};
25
26/// Rewrite column references in a physical expr according to a mapping.
27///
28/// This rewriter traverses the expression tree and replaces [`Column`] nodes
29/// with the corresponding expression found in the `column_map`.
30///
31/// If a column is found in the map, it is replaced by the mapped expression.
32/// If a column is NOT found in the map, a `DataFusionError::Internal` is
33/// returned.
34pub struct PhysicalColumnRewriter<'a> {
35    /// Mapping from original column to new column.
36    pub column_map: &'a HashMap<Column, Arc<dyn PhysicalExpr>>,
37}
38
39impl<'a> PhysicalColumnRewriter<'a> {
40    /// Create a new PhysicalColumnRewriter with the given column mapping.
41    pub fn new(column_map: &'a HashMap<Column, Arc<dyn PhysicalExpr>>) -> Self {
42        Self { column_map }
43    }
44}
45
46impl<'a> TreeNodeRewriter for PhysicalColumnRewriter<'a> {
47    type Node = Arc<dyn PhysicalExpr>;
48
49    fn f_down(
50        &mut self,
51        node: Self::Node,
52    ) -> datafusion_common::Result<Transformed<Self::Node>> {
53        if let Some(column) = node.downcast_ref::<Column>() {
54            if let Some(new_column) = self.column_map.get(column) {
55                // jump to prevent rewriting the new sub-expression again
56                return Ok(Transformed::new(
57                    Arc::clone(new_column),
58                    true,
59                    TreeNodeRecursion::Jump,
60                ));
61            } else {
62                // Column not found in mapping
63                return Err(DataFusionError::Internal(format!(
64                    "Column {column:?} not found in column mapping {:?}",
65                    self.column_map
66                )));
67            }
68        }
69        Ok(Transformed::no(node))
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76    use arrow::datatypes::{DataType, Field, Schema};
77    use datafusion_common::{Result, tree_node::TreeNode};
78    use datafusion_physical_expr::{
79        PhysicalExpr,
80        expressions::{Column, binary, col, lit},
81    };
82
83    /// Helper function to create a test schema
84    fn create_test_schema() -> Arc<Schema> {
85        Arc::new(Schema::new(vec![
86            Field::new("a", DataType::Int32, true),
87            Field::new("b", DataType::Int32, true),
88            Field::new("c", DataType::Int32, true),
89            Field::new("d", DataType::Int32, true),
90            Field::new("e", DataType::Int32, true),
91            Field::new("new_col", DataType::Int32, true),
92            Field::new("inner_col", DataType::Int32, true),
93            Field::new("another_col", DataType::Int32, true),
94        ]))
95    }
96
97    /// Helper function to create a complex nested expression with multiple columns
98    /// Create: (col_a + col_b) * (col_c - col_d) + col_e
99    fn create_complex_expression(schema: &Schema) -> Arc<dyn PhysicalExpr> {
100        let col_a = col("a", schema).unwrap();
101        let col_b = col("b", schema).unwrap();
102        let col_c = col("c", schema).unwrap();
103        let col_d = col("d", schema).unwrap();
104        let col_e = col("e", schema).unwrap();
105
106        let add_expr =
107            binary(col_a, datafusion_expr::Operator::Plus, col_b, schema).unwrap();
108        let sub_expr =
109            binary(col_c, datafusion_expr::Operator::Minus, col_d, schema).unwrap();
110        let mul_expr = binary(
111            add_expr,
112            datafusion_expr::Operator::Multiply,
113            sub_expr,
114            schema,
115        )
116        .unwrap();
117        binary(mul_expr, datafusion_expr::Operator::Plus, col_e, schema).unwrap()
118    }
119
120    /// Helper function to create a deeply nested expression
121    /// Create: col_a + (col_b + (col_c + (col_d + col_e)))
122    fn create_deeply_nested_expression(schema: &Schema) -> Arc<dyn PhysicalExpr> {
123        let col_a = col("a", schema).unwrap();
124        let col_b = col("b", schema).unwrap();
125        let col_c = col("c", schema).unwrap();
126        let col_d = col("d", schema).unwrap();
127        let col_e = col("e", schema).unwrap();
128
129        let inner1 =
130            binary(col_d, datafusion_expr::Operator::Plus, col_e, schema).unwrap();
131        let inner2 =
132            binary(col_c, datafusion_expr::Operator::Plus, inner1, schema).unwrap();
133        let inner3 =
134            binary(col_b, datafusion_expr::Operator::Plus, inner2, schema).unwrap();
135        binary(col_a, datafusion_expr::Operator::Plus, inner3, schema).unwrap()
136    }
137
138    #[test]
139    fn test_simple_column_replacement_with_jump() -> Result<()> {
140        let schema = create_test_schema();
141
142        // Test that Jump prevents re-processing of replaced columns
143        let mut column_map = HashMap::new();
144        column_map.insert(Column::new_with_schema("a", &schema).unwrap(), lit(42i32));
145        column_map.insert(
146            Column::new_with_schema("b", &schema).unwrap(),
147            lit("replaced_b"),
148        );
149        column_map.insert(
150            Column::new_with_schema("c", &schema).unwrap(),
151            col("c", &schema).unwrap(),
152        );
153        column_map.insert(
154            Column::new_with_schema("d", &schema).unwrap(),
155            col("d", &schema).unwrap(),
156        );
157        column_map.insert(
158            Column::new_with_schema("e", &schema).unwrap(),
159            col("e", &schema).unwrap(),
160        );
161
162        let mut rewriter = PhysicalColumnRewriter::new(&column_map);
163        let expr = create_complex_expression(&schema);
164
165        let result = expr.rewrite(&mut rewriter)?;
166
167        // Verify the transformation occurred
168        assert!(result.transformed);
169
170        assert_eq!(
171            format!("{}", result.data),
172            "(42 + replaced_b) * (c@2 - d@3) + e@4"
173        );
174
175        Ok(())
176    }
177
178    #[test]
179    fn test_nested_column_replacement_with_jump() -> Result<()> {
180        let schema = create_test_schema();
181        // Test Jump behavior with deeply nested expressions
182        let mut column_map = HashMap::new();
183        // Replace col_c with a complex expression containing new columns
184        let replacement_expr = binary(
185            lit(100i32),
186            datafusion_expr::Operator::Plus,
187            col("new_col", &schema).unwrap(),
188            &schema,
189        )
190        .unwrap();
191        column_map.insert(
192            Column::new_with_schema("c", &schema).unwrap(),
193            replacement_expr,
194        );
195        column_map.insert(
196            Column::new_with_schema("a", &schema).unwrap(),
197            col("a", &schema).unwrap(),
198        );
199        column_map.insert(
200            Column::new_with_schema("b", &schema).unwrap(),
201            col("b", &schema).unwrap(),
202        );
203        column_map.insert(
204            Column::new_with_schema("d", &schema).unwrap(),
205            col("d", &schema).unwrap(),
206        );
207        column_map.insert(
208            Column::new_with_schema("e", &schema).unwrap(),
209            col("e", &schema).unwrap(),
210        );
211
212        let mut rewriter = PhysicalColumnRewriter::new(&column_map);
213        let expr = create_deeply_nested_expression(&schema);
214
215        let result = expr.rewrite(&mut rewriter)?;
216
217        // Verify transformation occurred
218        assert!(result.transformed);
219
220        assert_eq!(
221            format!("{}", result.data),
222            "a@0 + b@1 + 100 + new_col@5 + d@3 + e@4"
223        );
224
225        Ok(())
226    }
227
228    #[test]
229    fn test_circular_reference_prevention() -> Result<()> {
230        let schema = create_test_schema();
231        // Test that Jump prevents infinite recursion with circular references
232        let mut column_map = HashMap::new();
233
234        // Create a circular reference: col_a -> col_b -> col_a (but Jump should prevent the second visit)
235        column_map.insert(
236            Column::new_with_schema("a", &schema).unwrap(),
237            col("b", &schema).unwrap(),
238        );
239        column_map.insert(
240            Column::new_with_schema("b", &schema).unwrap(),
241            col("a", &schema).unwrap(),
242        );
243
244        let mut rewriter = PhysicalColumnRewriter::new(&column_map);
245
246        // Start with an expression containing col_a
247        let expr = binary(
248            col("a", &schema).unwrap(),
249            datafusion_expr::Operator::Plus,
250            col("b", &schema).unwrap(),
251            &schema,
252        )
253        .unwrap();
254
255        let result = expr.rewrite(&mut rewriter)?;
256
257        // Verify transformation occurred
258        assert!(result.transformed);
259
260        assert_eq!(format!("{}", result.data), "b@1 + a@0");
261
262        Ok(())
263    }
264
265    #[test]
266    fn test_multiple_replacements_in_same_expression() -> Result<()> {
267        let schema = create_test_schema();
268        // Test multiple column replacements in the same complex expression
269        let mut column_map = HashMap::new();
270
271        // Replace multiple columns with literals
272        column_map.insert(Column::new_with_schema("a", &schema).unwrap(), lit(10i32));
273        column_map.insert(Column::new_with_schema("c", &schema).unwrap(), lit(20i32));
274        column_map.insert(Column::new_with_schema("e", &schema).unwrap(), lit(30i32));
275        column_map.insert(
276            Column::new_with_schema("b", &schema).unwrap(),
277            col("b", &schema).unwrap(),
278        );
279        column_map.insert(
280            Column::new_with_schema("d", &schema).unwrap(),
281            col("d", &schema).unwrap(),
282        );
283
284        let mut rewriter = PhysicalColumnRewriter::new(&column_map);
285        let expr = create_complex_expression(&schema); // (col_a + col_b) * (col_c - col_d) + col_e
286
287        let result = expr.rewrite(&mut rewriter)?;
288
289        // Verify transformation occurred
290        assert!(result.transformed);
291        assert_eq!(format!("{}", result.data), "(10 + b@1) * (20 - d@3) + 30");
292
293        Ok(())
294    }
295
296    #[test]
297    fn test_jump_with_complex_replacement_expression() -> Result<()> {
298        let schema = create_test_schema();
299        // Test Jump behavior when replacing with very complex expressions
300        let mut column_map = HashMap::new();
301
302        // Replace col_a with a complex nested expression
303        let inner_expr = binary(
304            lit(5i32),
305            datafusion_expr::Operator::Multiply,
306            col("a", &schema).unwrap(),
307            &schema,
308        )
309        .unwrap();
310        let middle_expr = binary(
311            inner_expr,
312            datafusion_expr::Operator::Plus,
313            lit(3i32),
314            &schema,
315        )
316        .unwrap();
317        let complex_replacement = binary(
318            middle_expr,
319            datafusion_expr::Operator::Minus,
320            col("another_col", &schema).unwrap(),
321            &schema,
322        )
323        .unwrap();
324
325        column_map.insert(
326            Column::new_with_schema("a", &schema).unwrap(),
327            complex_replacement,
328        );
329        column_map.insert(
330            Column::new_with_schema("b", &schema).unwrap(),
331            col("b", &schema).unwrap(),
332        );
333
334        let mut rewriter = PhysicalColumnRewriter::new(&column_map);
335
336        // Create expression: col_a + col_b
337        let expr = binary(
338            col("a", &schema).unwrap(),
339            datafusion_expr::Operator::Plus,
340            col("b", &schema).unwrap(),
341            &schema,
342        )
343        .unwrap();
344
345        let result = expr.rewrite(&mut rewriter)?;
346
347        assert_eq!(
348            format!("{}", result.data),
349            "5 * a@0 + 3 - another_col@7 + b@1"
350        );
351
352        // Verify transformation occurred
353        assert!(result.transformed);
354
355        Ok(())
356    }
357
358    #[test]
359    fn test_unmapped_columns_detection() -> Result<()> {
360        let schema = create_test_schema();
361        let mut column_map = HashMap::new();
362
363        // Only map col_a, leave col_b unmapped
364        column_map.insert(Column::new_with_schema("a", &schema).unwrap(), lit(42i32));
365
366        let mut rewriter = PhysicalColumnRewriter::new(&column_map);
367
368        // Create expression: col_a + col_b
369        let expr = binary(
370            col("a", &schema).unwrap(),
371            datafusion_expr::Operator::Plus,
372            col("b", &schema).unwrap(),
373            &schema,
374        )
375        .unwrap();
376
377        let err = expr.rewrite(&mut rewriter).unwrap_err();
378        assert!(matches!(err, DataFusionError::Internal(_)));
379
380        Ok(())
381    }
382}
datafusion_physical_plan/column_rewriter.rs

datafusion_physical_plan/
column_rewriter.rs
