Skip to main content

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::Expression;
12use crate::optimizer::ArrayOptimizer;
13use crate::scalar_fn::fns::literal::Literal;
14use crate::scalar_fn::fns::root::Root;
15
16impl ArrayRef {
17    /// Apply the expression to this array, producing a new array in constant time.
18    pub fn apply(self, expr: &Expression) -> VortexResult<ArrayRef> {
19        // If the expression is a root, return self.
20        if expr.is::<Root>() {
21            return Ok(self);
22        }
23
24        // Manually convert literals to ConstantArray.
25        if let Some(scalar) = expr.as_opt::<Literal>() {
26            return Ok(ConstantArray::new(scalar.clone(), self.len()).into_array());
27        }
28
29        // Otherwise, collect the child arrays.
30        let children: Vec<_> = expr
31            .children()
32            .iter()
33            .map(|e| self.clone().apply(e))
34            .try_collect()?;
35
36        // And wrap the scalar function up in an array.
37        let array =
38            ScalarFnArray::try_new(expr.scalar_fn().clone(), children, self.len())?.into_array();
39
40        // Optimize the resulting array's root.
41        array.optimize()
42    }
43}