vortex_array/expr/
optimize.rs

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

vortex_array/expr/
optimize.rs
