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::VortexResult;
6
7use crate::ArrayRef;
8use crate::IntoArray;
9use crate::arrays::ConstantArray;
10use crate::arrays::ScalarFnArray;
11use crate::expr::BoundExpression;
12use crate::expr::BoundKind;
13use crate::expr::Expression;
14use crate::optimizer::ArrayOptimizer;
15use crate::scalar_fn::fns::literal::Literal;
16use crate::scalar_fn::fns::root::Root;
17
18impl ArrayRef {
19    /// Apply a bound expression to this array, producing a new array in constant time.
20    pub fn apply_bound(self, expr: &BoundExpression) -> VortexResult<ArrayRef> {
21        let BoundKind::Scalar {
22            scalar_fn,
23            children,
24        } = expr.kind()
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 array =
65            ScalarFnArray::try_new_with_len(expr.scalar_fn().clone(), children, self.len())?
66                .into_array();
67
68        // Optimize the resulting array's root.
69        array.optimize()
70    }
71}
vortex_array/expression.rs

vortex_array/
expression.rs
