-
Notifications
You must be signed in to change notification settings - Fork 203
Expand file tree
/
Copy pathreplace.rs
More file actions
74 lines (67 loc) · 2.29 KB
/
Copy pathreplace.rs
File metadata and controls
74 lines (67 loc) · 2.29 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors
use vortex_error::VortexExpect;
use crate::dtype::Nullability;
use crate::dtype::StructFields;
use crate::expr::Expression;
use crate::expr::col;
use crate::expr::pack;
use crate::expr::root;
use crate::expr::traversal::NodeExt;
use crate::expr::traversal::Transformed;
use crate::expr::traversal::TraversalOrder;
/// Replaces all occurrences of `needle` in the expression `expr` with `replacement`.
pub fn replace(expr: Expression, needle: &Expression, replacement: Expression) -> Expression {
expr.transform_down(|node| {
if &node == needle {
Ok(Transformed {
value: replacement.clone(),
// If there is a match with a needle there can be no more matches in that subtree.
order: TraversalOrder::Skip,
changed: true,
})
} else {
Ok(Transformed::no(node))
}
})
.vortex_expect("ReplaceVisitor should not fail")
.into_inner()
}
/// Expand the `root` expression with a pack of the given struct fields.
pub fn replace_root_fields(expr: Expression, fields: &StructFields) -> Expression {
replace(
expr,
&root(),
pack(
fields
.names()
.iter()
.map(|name| (name.clone(), col(name.clone()))),
Nullability::NonNullable,
),
)
}
#[cfg(test)]
mod test {
use super::replace;
use crate::dtype::Nullability::NonNullable;
use crate::expr::get_item;
use crate::expr::lit;
use crate::expr::pack;
#[test]
fn test_replace_full_tree() {
let e = get_item("b", pack([("a", lit(1)), ("b", lit(2))], NonNullable));
let needle = get_item("b", pack([("a", lit(1)), ("b", lit(2))], NonNullable));
let replacement = lit(42);
let replaced_expr = replace(e, &needle, replacement.clone());
assert_eq!(&replaced_expr, &replacement);
}
#[test]
fn test_replace_leaf() {
let e = pack([("a", lit(1)), ("b", lit(2))], NonNullable);
let needle = lit(2);
let replacement = lit(42);
let replaced_expr = replace(e, &needle, replacement);
assert_eq!(replaced_expr.to_string(), "pack(a: 1i32, b: 42i32)");
}
}