vortex_array/
expression.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use smallvec::SmallVec;
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 len = self.len();
34        let children = children
35            .iter()
36            .map(|child| self.clone().apply_bound(child))
37            .collect::<VortexResult<SmallVec<[ArrayRef; 4]>>>()?;
38
39        let array = ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, 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 len = self.len();
58        let children = expr
59            .children()
60            .iter()
61            .map(|e| self.clone().apply(e))
62            .collect::<VortexResult<SmallVec<[ArrayRef; 4]>>>()?;
63
64        // And wrap the scalar function up in an array.
65        let scalar_fn = expr
66            .as_scalar()
67            .vortex_expect("root and literal were handled above, so this is a scalar node");
68        let array = ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, len)?.into_array();
69
70        // Optimize the resulting array's root.
71        array.optimize()
72    }
73}
vortex_array/expression.rs

vortex_array/
expression.rs
