vortex_array/arrays/scalar_fn/vtable/
validity.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5
6use crate::ArrayRef;
7use crate::IntoArray;
8use crate::LEGACY_SESSION;
9use crate::VortexSessionExecute;
10use crate::arrays::scalar_fn::array::ScalarFnArray;
11use crate::arrays::scalar_fn::vtable::ArrayExpr;
12use crate::arrays::scalar_fn::vtable::FakeEq;
13use crate::arrays::scalar_fn::vtable::ScalarFnVTable;
14use crate::expr::ExecutionArgs;
15use crate::expr::Expression;
16use crate::expr::Literal;
17use crate::expr::Root;
18use crate::expr::ScalarFn;
19use crate::expr::lit;
20use crate::validity::Validity;
21use crate::vtable::ValidityVTable;
22
23/// Execute an expression tree recursively.
24///
25/// This assumes all leaf expressions are either ArrayExpr (wrapping actual arrays) or Literals.
26fn execute_expr(expr: &Expression, row_count: usize) -> VortexResult<ArrayRef> {
27    let mut ctx = LEGACY_SESSION.create_execution_ctx();
28
29    // Handle Root expression - this should not happen in validity expressions
30    if expr.is::<Root>() {
31        vortex_error::vortex_bail!("Root expression cannot be executed in validity context");
32    }
33
34    // Handle Literal expression - create a constant array
35    if expr.is::<Literal>() {
36        let scalar = expr.as_::<Literal>();
37        return Ok(crate::arrays::ConstantArray::new(scalar.clone(), row_count).into_array());
38    }
39
40    // Recursively execute child expressions to get input arrays
41    let inputs: Vec<ArrayRef> = expr
42        .children()
43        .iter()
44        .map(|child| execute_expr(child, row_count))
45        .collect::<VortexResult<_>>()?;
46
47    let args = ExecutionArgs {
48        inputs,
49        row_count,
50        ctx: &mut ctx,
51    };
52
53    Ok(expr.scalar_fn().execute(args)?.into_array())
54}
55
56impl ValidityVTable<ScalarFnVTable> for ScalarFnVTable {
57    fn validity(array: &ScalarFnArray) -> VortexResult<Validity> {
58        let inputs: Vec<_> = array
59            .children
60            .iter()
61            .map(|child| {
62                if let Some(scalar) = child.as_constant() {
63                    return Ok(lit(scalar));
64                }
65                Expression::try_new(ScalarFn::new(ArrayExpr, FakeEq(child.clone())), [])
66            })
67            .collect::<VortexResult<_>>()?;
68
69        let expr = Expression::try_new(array.scalar_fn.clone(), inputs)?;
70        let validity_expr = array.scalar_fn().validity(&expr)?;
71
72        // Execute the validity expression. All leaves are ArrayExpr nodes.
73        Ok(Validity::Array(execute_expr(&validity_expr, array.len())?))
74    }
75}
vortex_array/arrays/scalar_fn/vtable/validity.rs

vortex_array/arrays/scalar_fn/vtable/
validity.rs
