vortex_array/arrays/primitive/vtable/
operator.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_buffer::Buffer;
5use vortex_compute::filter::Filter;
6use vortex_dtype::match_each_native_ptype;
7use vortex_error::VortexResult;
8use vortex_vector::primitive::PVector;
9
10use crate::arrays::{MaskedVTable, PrimitiveArray, PrimitiveVTable};
11use crate::execution::{BatchKernelRef, BindCtx, kernel};
12use crate::vtable::{OperatorVTable, ValidityHelper};
13use crate::{ArrayRef, IntoArray};
14
15impl OperatorVTable<PrimitiveVTable> for PrimitiveVTable {
16    fn bind(
17        array: &PrimitiveArray,
18        selection: Option<&ArrayRef>,
19        ctx: &mut dyn BindCtx,
20    ) -> VortexResult<BatchKernelRef> {
21        let mask = ctx.bind_selection(array.len(), selection)?;
22        let validity = ctx.bind_validity(array.validity(), array.len(), selection)?;
23
24        match_each_native_ptype!(array.ptype(), |P| {
25            let elements = array.buffer::<P>();
26            Ok(kernel(move || {
27                let mask = mask.execute()?;
28                let validity = validity.execute()?;
29
30                // Note that validity already has the mask applied so we only need to apply it to
31                // the elements.
32                let elements = elements.filter(&mask);
33
34                Ok(PVector::<P>::try_new(elements, validity)?.into())
35            }))
36        })
37    }
38
39    fn reduce_parent(
40        array: &PrimitiveArray,
41        parent: &ArrayRef,
42        _child_idx: usize,
43    ) -> VortexResult<Option<ArrayRef>> {
44        // Push-down masking of `validity` from the parent `MaskedArray`.
45        if let Some(masked) = parent.as_opt::<MaskedVTable>() {
46            let masked_array = match_each_native_ptype!(array.ptype(), |T| {
47                // SAFETY: Since we are only flipping some bits in the validity, all invariants that
48                // were upheld are still upheld.
49                unsafe {
50                    PrimitiveArray::new_unchecked(
51                        Buffer::<T>::from_byte_buffer(array.byte_buffer().clone()),
52                        array.validity().clone().and(masked.validity().clone()),
53                    )
54                }
55                .into_array()
56            });
57
58            return Ok(Some(masked_array));
59        }
60
61        Ok(None)
62    }
63}