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