vortex_array/expr/transform/
simplify.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5
6use crate::expr::Expression;
7use crate::expr::exprs::get_item::GetItem;
8use crate::expr::exprs::pack::Pack;
9use crate::expr::transform::match_between::find_between;
10use crate::expr::traversal::{NodeExt, Transformed};
11
12/// Simplifies an expression into an equivalent expression which is faster and easier to analyze.
13///
14/// If the scope dtype is known, see `simplify_typed` for a simplifier which uses dtype.
15pub fn simplify(e: Expression) -> VortexResult<Expression> {
16    let e = e
17        .transform_up(simplify_transformer)
18        .map(|e| e.into_inner())?;
19    Ok(find_between(e))
20}
21
22fn simplify_transformer(node: Expression) -> VortexResult<Transformed<Expression>> {
23    // pack(l_1: e_1, ..., l_i: e_i, ..., l_n: e_n).get_item(l_i) = e_i where 0 <= i <= n
24    if let Some(get_item) = node.as_opt::<GetItem>()
25        && let Some(pack) = get_item.child(0).as_opt::<Pack>()
26    {
27        let expr = pack.field(get_item.data())?;
28        return Ok(Transformed::yes(expr));
29    }
30    Ok(Transformed::no(node))
31}
32
33#[cfg(test)]
34mod tests {
35    use vortex_dtype::Nullability::NonNullable;
36
37    use super::simplify;
38    use crate::expr::exprs::get_item::get_item;
39    use crate::expr::exprs::literal::lit;
40    use crate::expr::exprs::pack::pack;
41
42    #[test]
43    fn test_simplify() {
44        let e = get_item("b", pack([("a", lit(1)), ("b", lit(2))], NonNullable));
45        let e = simplify(e).unwrap();
46        assert_eq!(&e, &lit(2));
47    }
48}