vortex-array 0.78.0

Vortex in memory columnar data format
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use vortex_error::VortexResult;
use vortex_error::vortex_bail;

use crate::ArrayRef;
use crate::ExecutionCtx;
use crate::IntoArray;
use crate::VortexSessionExecute;
use crate::array::ArrayView;
use crate::array::ValidityVTable;
use crate::arrays::ConstantArray;
use crate::arrays::scalar_fn::ScalarFnArrayExt;
use crate::arrays::scalar_fn::vtable::ArrayExpr;
use crate::arrays::scalar_fn::vtable::FakeEq;
use crate::arrays::scalar_fn::vtable::ScalarFn;
use crate::expr::Expression;
use crate::expr::lit;
use crate::legacy_session;
use crate::scalar_fn::TypedScalarFnInstance;
use crate::scalar_fn::VecExecutionArgs;
use crate::scalar_fn::fns::literal::Literal;
use crate::scalar_fn::fns::root::Root;
use crate::validity::Validity;

/// Execute an expression tree recursively.
///
/// This assumes all leaf expressions are either ArrayExpr (wrapping actual arrays) or Literals.
fn execute_expr(
    expr: &Expression,
    row_count: usize,
    ctx: &mut ExecutionCtx,
) -> VortexResult<ArrayRef> {
    // Handle Root expression - this should not happen in validity expressions
    if expr.is::<Root>() {
        vortex_bail!("Root expression cannot be executed in validity context");
    }

    // Handle Literal expression - create a constant array
    if expr.is::<Literal>() {
        let scalar = expr.as_::<Literal>();
        return Ok(ConstantArray::new(scalar.clone(), row_count).into_array());
    }

    // Recursively execute child expressions to get input arrays
    let inputs: Vec<ArrayRef> = expr
        .children()
        .iter()
        .map(|child| execute_expr(child, row_count, ctx))
        .collect::<VortexResult<_>>()?;

    let args = VecExecutionArgs::new(inputs, row_count);

    Ok(expr.scalar_fn().execute(&args, ctx)?.into_array())
}

impl ValidityVTable<ScalarFn> for ScalarFn {
    fn validity(array: ArrayView<'_, ScalarFn>) -> VortexResult<Validity> {
        let inputs: Vec<_> = array
            .iter_children()
            .map(|child| {
                if let Some(scalar) = child.as_constant() {
                    return Ok(lit(scalar));
                }
                Expression::try_new(
                    TypedScalarFnInstance::new(ArrayExpr, FakeEq(child.clone())).erased(),
                    [],
                )
            })
            .collect::<VortexResult<_>>()?;

        let expr = Expression::try_new(array.scalar_fn().clone(), inputs)?;
        let validity_expr = array.scalar_fn().validity(&expr)?;

        #[allow(clippy::disallowed_methods)]
        let ctx = &mut legacy_session().create_execution_ctx();
        // Execute the validity expression. All leaves are ArrayExpr nodes.
        Ok(Validity::Array(execute_expr(
            &validity_expr,
            array.len(),
            ctx,
        )?))
    }
}