vortex-array 0.86.1

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

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

use vortex_error::VortexResult;
use vortex_mask::Mask;
use vortex_mask::MaskValuesRef;

use super::super::RowFnExecutionArgs;
use super::super::args::BorrowedRowFnArgs;
use crate::ArrayRef;
use crate::ExecutionCtx;
use crate::IntoArray;
use crate::builtins::ArrayBuiltins;

impl RowFnExecutionArgs {
    /// Resolve validity and try direct valid-row execution before filtered execution.
    pub(super) fn execute_valid_only(
        &self,
        kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult<ArrayRef>,
        try_valid_rows: impl FnOnce(
            BorrowedRowFnArgs<'_>,
            MaskValuesRef,
            &mut ExecutionCtx,
        ) -> VortexResult<Option<ArrayRef>>,
        execute_filtered_rows: impl FnOnce(
            BorrowedRowFnArgs<'_>,
            MaskValuesRef,
            &mut ExecutionCtx,
        ) -> VortexResult<ArrayRef>,
        ctx: &mut ExecutionCtx,
    ) -> VortexResult<ArrayRef> {
        let validity = self.validity.clone().execute_mask(self.row_count, ctx)?;

        let valid_rows = match validity {
            // An empty mask is both all-valid and all-null. Preserve the all-valid behavior.
            Mask::AllTrue(_) | Mask::AllFalse(0) => {
                let values = kernel(self.execution_args(&self.inputs, self.row_count), ctx)?;
                let values = self.validate_kernel_output(values, self.row_count, ctx)?;

                return self.finalize_output(values, self.row_count);
            }
            Mask::AllFalse(_) => return Ok(self.all_null()),
            Mask::Values(valid_rows) => valid_rows,
        };

        if let Some(result) = self.try_execute_valid_rows(try_valid_rows, &valid_rows, ctx)? {
            return Ok(result);
        }

        self.execute_filtered(execute_filtered_rows, &valid_rows, ctx)
    }

    /// Try execution against the original inputs, then mask a returned full-length result.
    pub(super) fn try_execute_valid_rows(
        &self,
        try_valid_rows: impl FnOnce(
            BorrowedRowFnArgs<'_>,
            MaskValuesRef,
            &mut ExecutionCtx,
        ) -> VortexResult<Option<ArrayRef>>,
        valid: &MaskValuesRef,
        ctx: &mut ExecutionCtx,
    ) -> VortexResult<Option<ArrayRef>> {
        let Some(values) = try_valid_rows(
            self.execution_args(&self.inputs, self.row_count),
            MaskValuesRef::clone(valid),
            ctx,
        )?
        else {
            return Ok(None);
        };
        let values = self.validate_kernel_output(values, valid.len(), ctx)?;

        let mask = valid.as_ref().into_array();
        self.finalize_output(values.mask(mask)?, valid.len())
            .map(Some)
    }
}