vortex_array/expr/transform/
replace.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexExpect;
5
6use crate::expr::Expression;
7use crate::expr::traversal::NodeExt;
8use crate::expr::traversal::Transformed;
9use crate::expr::traversal::TraversalOrder;
10
11/// Replaces all occurrences of `needle` in the expression `expr` with `replacement`.
12pub fn replace(expr: Expression, needle: &Expression, replacement: Expression) -> Expression {
13    expr.transform_down(|node| {
14        if &node == needle {
15            Ok(Transformed {
16                value: replacement.clone(),
17                // If there is a match with a needle there can be no more matches in that subtree.
18                order: TraversalOrder::Skip,
19                changed: true,
20            })
21        } else {
22            Ok(Transformed::no(node))
23        }
24    })
25    .vortex_expect("ReplaceVisitor should not fail")
26    .into_inner()
27}
28
29#[cfg(test)]
30mod test {
31    use super::replace;
32    use crate::dtype::Nullability::NonNullable;
33    use crate::expr::get_item;
34    use crate::expr::lit;
35    use crate::expr::pack;
36
37    #[test]
38    fn test_replace_full_tree() {
39        let e = get_item("b", pack([("a", lit(1)), ("b", lit(2))], NonNullable));
40        let needle = get_item("b", pack([("a", lit(1)), ("b", lit(2))], NonNullable));
41        let replacement = lit(42);
42        let replaced_expr = replace(e, &needle, replacement.clone());
43        assert_eq!(&replaced_expr, &replacement);
44    }
45
46    #[test]
47    fn test_replace_leaf() {
48        let e = pack([("a", lit(1)), ("b", lit(2))], NonNullable);
49        let needle = lit(2);
50        let replacement = lit(42);
51        let replaced_expr = replace(e, &needle, replacement);
52        assert_eq!(replaced_expr.to_string(), "pack(a: 1i32, b: 42i32)");
53    }
54}
vortex_array/expr/transform/replace.rs

vortex_array/expr/transform/
replace.rs
