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

vortex_array/arrays/scalar_fn/vtable/
validity.rs
