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;
5use vortex_error::vortex_bail;
6
7use crate::ArrayRef;
8use crate::ExecutionCtx;
9use crate::IntoArray;
10use crate::VortexSessionExecute;
11use crate::array::ArrayView;
12use crate::array::ValidityVTable;
13use crate::arrays::ConstantArray;
14use crate::arrays::scalar_fn::ScalarFnArrayExt;
15use crate::arrays::scalar_fn::vtable::ArrayExpr;
16use crate::arrays::scalar_fn::vtable::FakeEq;
17use crate::arrays::scalar_fn::vtable::ScalarFn;
18use crate::expr::Expression;
19use crate::expr::lit;
20use crate::legacy_session;
21use crate::scalar_fn::TypedScalarFnInstance;
22use crate::scalar_fn::VecExecutionArgs;
23use crate::scalar_fn::fns::literal::Literal;
24use crate::validity::Validity;
25
26/// Execute an expression tree recursively.
27///
28/// This assumes all leaf expressions are either ArrayExpr (wrapping actual arrays) or Literals.
29fn execute_expr(
30    expr: &Expression,
31    row_count: usize,
32    ctx: &mut ExecutionCtx,
33) -> VortexResult<ArrayRef> {
34    // Only Expression::Scalar is executable
35    let Some(scalar_fn) = expr.as_scalar() else {
36        vortex_bail!("Only Expression::Scalar is executable");
37    };
38
39    // Handle Literal expression - create a constant array
40    if expr.is::<Literal>() {
41        let scalar = expr.as_::<Literal>();
42        return Ok(ConstantArray::new(scalar.clone(), row_count).into_array());
43    }
44
45    // Recursively execute child expressions to get input arrays
46    let inputs: Vec<ArrayRef> = expr
47        .children()
48        .iter()
49        .map(|child| execute_expr(child, row_count, ctx))
50        .collect::<VortexResult<_>>()?;
51
52    let args = VecExecutionArgs::new(inputs, row_count);
53
54    Ok(scalar_fn.execute(&args, ctx)?.into_array())
55}
56
57impl ValidityVTable<ScalarFn> for ScalarFn {
58    fn validity(array: ArrayView<'_, ScalarFn>) -> VortexResult<Validity> {
59        let inputs: Vec<_> = array
60            .iter_children()
61            .map(|child| {
62                if let Some(scalar) = child.as_constant() {
63                    return Ok(lit(scalar));
64                }
65                Expression::try_new(
66                    TypedScalarFnInstance::new(ArrayExpr, FakeEq(child.clone())).erased(),
67                    [],
68                )
69            })
70            .collect::<VortexResult<_>>()?;
71
72        let expr = Expression::try_new(array.scalar_fn().clone(), inputs)?;
73        let validity_expr = array.scalar_fn().validity(&expr)?;
74
75        #[allow(clippy::disallowed_methods)]
76        let ctx = &mut legacy_session().create_execution_ctx();
77        // Execute the validity expression. All leaves are ArrayExpr nodes.
78        Ok(Validity::Array(execute_expr(
79            &validity_expr,
80            array.len(),
81            ctx,
82        )?))
83    }
84}
vortex_array/arrays/scalar_fn/vtable/validity.rs

vortex_array/arrays/scalar_fn/vtable/
validity.rs
