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::dtype::Nullability;
7use crate::dtype::StructFields;
8use crate::expr::Expression;
9use crate::expr::col;
10use crate::expr::pack;
11use crate::expr::root;
12use crate::expr::traversal::NodeExt;
13use crate::expr::traversal::Transformed;
14use crate::expr::traversal::TraversalOrder;
15
16/// Replaces all occurrences of `needle` in the expression `expr` with `replacement`.
17pub fn replace(expr: Expression, needle: &Expression, replacement: Expression) -> Expression {
18    expr.transform_down(|node| {
19        if &node == needle {
20            Ok(Transformed {
21                value: replacement.clone(),
22                // If there is a match with a needle there can be no more matches in that subtree.
23                order: TraversalOrder::Skip,
24                changed: true,
25            })
26        } else {
27            Ok(Transformed::no(node))
28        }
29    })
30    .vortex_expect("ReplaceVisitor should not fail")
31    .into_inner()
32}
33
34/// Expand the `root` expression with a pack of the given struct fields.
35pub fn replace_root_fields(expr: Expression, fields: &StructFields) -> Expression {
36    replace(
37        expr,
38        &root(),
39        pack(
40            fields
41                .names()
42                .iter()
43                .map(|name| (name.clone(), col(name.clone()))),
44            Nullability::NonNullable,
45        ),
46    )
47}
48
49#[cfg(test)]
50mod test {
51    use super::replace;
52    use crate::dtype::Nullability::NonNullable;
53    use crate::expr::get_item;
54    use crate::expr::lit;
55    use crate::expr::pack;
56
57    #[test]
58    fn test_replace_full_tree() {
59        let e = get_item("b", pack([("a", lit(1)), ("b", lit(2))], NonNullable));
60        let needle = get_item("b", pack([("a", lit(1)), ("b", lit(2))], NonNullable));
61        let replacement = lit(42);
62        let replaced_expr = replace(e, &needle, replacement.clone());
63        assert_eq!(&replaced_expr, &replacement);
64    }
65
66    #[test]
67    fn test_replace_leaf() {
68        let e = pack([("a", lit(1)), ("b", lit(2))], NonNullable);
69        let needle = lit(2);
70        let replacement = lit(42);
71        let replaced_expr = replace(e, &needle, replacement);
72        assert_eq!(replaced_expr.to_string(), "pack(a: 1i32, b: 42i32)");
73    }
74}
vortex_array/expr/transform/replace.rs

vortex_array/expr/transform/
replace.rs
