vortex_array/expr/
optimize.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5
6use crate::expr::BoundExpression;
7use crate::expr::transform::match_between::find_between;
8use crate::scalar_fn::ExpressionReduceNode;
9
10impl BoundExpression {
11    /// Optimize the root expression node only, iterating to convergence.
12    ///
13    /// This applies optimization rules repeatedly until no more changes occur:
14    /// 1. `simplify` - scalar-function simplifications over the bound node
15    /// 2. `reduce` - abstract reduction rules via `ReduceNode`
16    pub fn optimize(&self) -> VortexResult<BoundExpression> {
17        Ok(self.try_optimize()?.unwrap_or_else(|| self.clone()))
18    }
19
20    /// Apply this node's own simplification rule, if it has one.
21    ///
22    /// Non-scalar nodes carry no rules, so they never simplify.
23    fn simplify_node(&self) -> VortexResult<Option<BoundExpression>> {
24        match self {
25            BoundExpression::Scalar { scalar_fn, .. } => scalar_fn.simplify(self),
26            BoundExpression::Root { .. } => Ok(None),
27        }
28    }
29
30    /// Apply this node's own abstract reduction rule, if it has one.
31    fn reduce_node<'a>(
32        &self,
33        node: &ExpressionReduceNode<'a>,
34    ) -> VortexResult<Option<ExpressionReduceNode<'a>>> {
35        match self {
36            BoundExpression::Scalar { scalar_fn, .. } => scalar_fn.reduce_expression(node),
37            BoundExpression::Root { .. } => Ok(None),
38        }
39    }
40
41    /// Try to optimize the root expression node only, returning None if no optimizations applied.
42    fn try_optimize(&self) -> VortexResult<Option<BoundExpression>> {
43        // Copy-on-write: `current` stays None until a rule fires, so unchanged nodes (the common
44        // case) are never cloned.
45        let mut current: Option<BoundExpression> = None;
46        let mut loop_counter = 0;
47
48        loop {
49            if loop_counter > 100 {
50                vortex_error::vortex_bail!(
51                    "Exceeded maximum optimization iterations (possible infinite loop)"
52                );
53            }
54            loop_counter += 1;
55
56            let expr = current.as_ref().unwrap_or(self);
57            let mut changed = false;
58
59            if let Some(simplified) = expr.simplify_node()? {
60                current = Some(simplified);
61                changed = true;
62            }
63
64            // Try reduce via ReduceNode. The node borrows the expression, so constructing it is
65            // free; the block scopes the borrow so `current` can be updated.
66            let reduced = {
67                let expr = current.as_ref().unwrap_or(self);
68                expr.reduce_node(&ExpressionReduceNode::new(expr))?
69                    .map(ExpressionReduceNode::into_expression)
70            };
71            if let Some(reduced_expr) = reduced {
72                current = Some(reduced_expr);
73                changed = true;
74            }
75
76            if !changed {
77                break;
78            }
79        }
80
81        Ok(current)
82    }
83
84    /// Optimize the entire expression tree recursively.
85    ///
86    /// Optimizes children first (bottom-up), then optimizes the root.
87    pub fn optimize_recursive(&self) -> VortexResult<BoundExpression> {
88        Ok(self
89            .try_optimize_recursive()?
90            .unwrap_or_else(|| self.clone()))
91    }
92
93    /// Try to optimize the entire expression tree recursively.
94    pub fn try_optimize_recursive(&self) -> VortexResult<Option<BoundExpression>> {
95        let result = self.try_optimize_recursive_inner()?;
96
97        // Apply the between optimization once at the top level only.
98        // TODO(ngates): remove the "between" optimization, or rewrite it to not always convert
99        //  to CNF?
100        Ok(Some(find_between(result.unwrap_or_else(|| self.clone()))))
101    }
102
103    fn try_optimize_recursive_inner(&self) -> VortexResult<Option<BoundExpression>> {
104        // First optimize the root
105        let mut current = self.try_optimize()?;
106
107        // Then recursively optimize children. The new children vector is only allocated once a
108        // child actually changes, so fully-optimized subtrees cost no allocations.
109        let expr = current.as_ref().unwrap_or(self);
110        let children = expr.children();
111        let mut new_children: Option<Vec<BoundExpression>> = None;
112        for (idx, child) in children.iter().enumerate() {
113            if let Some(optimized) = child.try_optimize_recursive_inner()? {
114                new_children
115                    .get_or_insert_with(|| children[..idx].to_vec())
116                    .push(optimized);
117            } else if let Some(new_children) = new_children.as_mut() {
118                new_children.push(child.clone());
119            }
120        }
121
122        if let Some(new_children) = new_children {
123            let updated = expr.clone().with_children(new_children)?;
124
125            // After updating children, try to optimize root again
126            current = Some(updated.try_optimize()?.unwrap_or(updated));
127        }
128
129        Ok(current)
130    }
131}
132
133#[cfg(test)]
134mod tests {
135    use rstest::rstest;
136    use vortex_error::VortexResult;
137    use vortex_error::vortex_err;
138
139    use crate::dtype::DType;
140    use crate::dtype::Nullability;
141    use crate::dtype::PType;
142    use crate::dtype::StructFields;
143    use crate::expr::Expression;
144    use crate::expr::and;
145    use crate::expr::cast;
146    use crate::expr::eq;
147    use crate::expr::get_item;
148    use crate::expr::lit;
149    use crate::expr::lt_eq;
150    use crate::expr::or;
151    use crate::expr::root;
152    use crate::expr::zip_expr;
153    use crate::scalar::Scalar;
154    use crate::scalar_fn::fns::literal::Literal;
155
156    #[test]
157    fn optimize_or_chain_correctness() -> VortexResult<()> {
158        let expr = or(
159            eq(get_item("x", root()), lit(1i32)),
160            eq(get_item("x", root()), lit(2i32)),
161        );
162        let scope = DType::Struct(
163            StructFields::new(
164                ["x"].into(),
165                vec![DType::Primitive(PType::I32, Nullability::NonNullable)],
166            ),
167            Nullability::NonNullable,
168        );
169        let optimized = expr.bind(&scope)?.optimize_recursive()?;
170
171        let s = optimized.to_string();
172        assert!(s.contains("$.x"), "expected $.x in {s}");
173        assert!(s.contains("1i32") || s.contains('1'), "expected 1 in {s}");
174        assert!(s.contains("2i32") || s.contains('2'), "expected 2 in {s}");
175        Ok(())
176    }
177
178    #[test]
179    fn optimize_folds_cast_of_literal_in_comparison() -> VortexResult<()> {
180        let expr = lt_eq(
181            get_item("x", root()),
182            cast(
183                lit(3i32),
184                DType::Primitive(PType::F64, Nullability::NonNullable),
185            ),
186        );
187        let scope = DType::Struct(
188            StructFields::new(
189                ["x"].into(),
190                vec![DType::Primitive(PType::F64, Nullability::NonNullable)],
191            ),
192            Nullability::NonNullable,
193        );
194        let optimized = expr.bind(&scope)?.optimize_recursive()?;
195
196        // Prune rules pattern-match a bare Literal on the comparison RHS; a cast wrapper
197        // silently disables pruning.
198        let rhs = optimized
199            .child(1)
200            .as_opt::<Literal>()
201            .ok_or_else(|| vortex_err!("expected a bare literal RHS, got {optimized}"))?;
202        assert_eq!(rhs, &Scalar::primitive(3.0f64, Nullability::NonNullable));
203        Ok(())
204    }
205
206    #[rstest]
207    #[case::and_annihilator(and(root(), lit(false)), DType::Bool(Nullability::Nullable))]
208    #[case::or_annihilator(or(root(), lit(true)), DType::Bool(Nullability::Nullable))]
209    #[case::and_nullable_identity(
210        and(root(), lit(Scalar::from(Some(true)))),
211        DType::Bool(Nullability::NonNullable)
212    )]
213    #[case::or_nullable_identity(
214        or(root(), lit(Scalar::from(Some(false)))),
215        DType::Bool(Nullability::NonNullable)
216    )]
217    #[case::zip_true(
218        zip_expr(lit(true), lit(1i32), root()),
219        DType::Primitive(PType::I32, Nullability::Nullable)
220    )]
221    #[case::zip_false(
222        zip_expr(lit(false), root(), lit(1i32)),
223        DType::Primitive(PType::I32, Nullability::Nullable)
224    )]
225    fn literal_simplification_preserves_result_dtype(
226        #[case] expr: Expression,
227        #[case] scope: DType,
228    ) -> VortexResult<()> {
229        let original = expr.bind(&scope)?;
230        let optimized = expr.bind(&scope)?.optimize_recursive()?;
231        assert_eq!(optimized.dtype(), original.dtype());
232        Ok(())
233    }
234}
vortex_array/expr/optimize.rs

vortex_array/expr/
optimize.rs
