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 vortex_error::VortexResult;
136    use vortex_error::vortex_err;
137
138    use crate::dtype::DType;
139    use crate::dtype::Nullability;
140    use crate::dtype::PType;
141    use crate::dtype::StructFields;
142    use crate::expr::cast;
143    use crate::expr::eq;
144    use crate::expr::get_item;
145    use crate::expr::lit;
146    use crate::expr::lt_eq;
147    use crate::expr::or;
148    use crate::expr::root;
149    use crate::scalar::Scalar;
150    use crate::scalar_fn::fns::literal::Literal;
151
152    #[test]
153    fn optimize_or_chain_correctness() -> VortexResult<()> {
154        let expr = or(
155            eq(get_item("x", root()), lit(1i32)),
156            eq(get_item("x", root()), lit(2i32)),
157        );
158        let scope = DType::Struct(
159            StructFields::new(
160                ["x"].into(),
161                vec![DType::Primitive(PType::I32, Nullability::NonNullable)],
162            ),
163            Nullability::NonNullable,
164        );
165        let optimized = expr.bind(&scope)?.optimize_recursive()?;
166
167        let s = optimized.to_string();
168        assert!(s.contains("$.x"), "expected $.x in {s}");
169        assert!(s.contains("1i32") || s.contains('1'), "expected 1 in {s}");
170        assert!(s.contains("2i32") || s.contains('2'), "expected 2 in {s}");
171        Ok(())
172    }
173
174    #[test]
175    fn optimize_folds_cast_of_literal_in_comparison() -> VortexResult<()> {
176        let expr = lt_eq(
177            get_item("x", root()),
178            cast(
179                lit(3i32),
180                DType::Primitive(PType::F64, Nullability::NonNullable),
181            ),
182        );
183        let scope = DType::Struct(
184            StructFields::new(
185                ["x"].into(),
186                vec![DType::Primitive(PType::F64, Nullability::NonNullable)],
187            ),
188            Nullability::NonNullable,
189        );
190        let optimized = expr.bind(&scope)?.optimize_recursive()?;
191
192        // Prune rules pattern-match a bare Literal on the comparison RHS; a cast wrapper
193        // silently disables pruning.
194        let rhs = optimized
195            .child(1)
196            .as_opt::<Literal>()
197            .ok_or_else(|| vortex_err!("expected a bare literal RHS, got {optimized}"))?;
198        assert_eq!(rhs, &Scalar::primitive(3.0f64, Nullability::NonNullable));
199        Ok(())
200    }
201}
vortex_array/expr/optimize.rs

vortex_array/expr/
optimize.rs
