vortex_array/
expression.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use itertools::Itertools;
5use vortex_error::VortexExpect;
6use vortex_error::VortexResult;
7
8use crate::ArrayRef;
9use crate::IntoArray;
10use crate::arrays::ConstantArray;
11use crate::arrays::ScalarFnArray;
12use crate::expr::BoundExpression;
13use crate::expr::Expression;
14use crate::optimizer::ArrayOptimizer;
15use crate::scalar_fn::fns::literal::Literal;
16
17impl ArrayRef {
18    /// Apply a bound expression to this array, producing a new array in constant time.
19    pub fn apply_bound(self, expr: &BoundExpression) -> VortexResult<ArrayRef> {
20        let BoundExpression::Scalar {
21            scalar_fn,
22            children,
23            ..
24        } = expr
25        else {
26            return Ok(self);
27        };
28
29        if let Some(scalar) = scalar_fn.as_opt::<Literal>() {
30            return Ok(ConstantArray::new(scalar.clone(), self.len()).into_array());
31        }
32
33        let children: Vec<_> = children
34            .iter()
35            .map(|child| self.clone().apply_bound(child))
36            .try_collect()?;
37
38        let array =
39            ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, self.len())?.into_array();
40
41        array.optimize()
42    }
43
44    /// Apply the expression to this array, producing a new array in constant time.
45    pub fn apply(self, expr: &Expression) -> VortexResult<ArrayRef> {
46        // If the expression is a root, return self.
47        if expr.is_root() {
48            return Ok(self);
49        }
50
51        // Manually convert literals to ConstantArray.
52        if let Some(scalar) = expr.as_opt::<Literal>() {
53            return Ok(ConstantArray::new(scalar.clone(), self.len()).into_array());
54        }
55
56        // Otherwise, collect the child arrays.
57        let children: Vec<_> = expr
58            .children()
59            .iter()
60            .map(|e| self.clone().apply(e))
61            .try_collect()?;
62
63        // And wrap the scalar function up in an array.
64        let scalar_fn = expr
65            .as_scalar()
66            .vortex_expect("root and literal were handled above, so this is a scalar node");
67        let array =
68            ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, self.len())?.into_array();
69
70        // Optimize the resulting array's root.
71        array.optimize()
72    }
73}
vortex_array/expression.rs

vortex_array/
expression.rs
