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_dtype::match_each_native_ptype;
6use vortex_error::VortexResult;
7
8use crate::ArrayRef;
9use crate::IntoArray;
10use crate::arrays::MaskedArray;
11use crate::arrays::MaskedVTable;
12use crate::arrays::PrimitiveArray;
13use crate::arrays::PrimitiveVTable;
14use crate::optimizer::rules::ArrayParentReduceRule;
15use crate::optimizer::rules::Exact;
16use crate::vtable::ValidityHelper;
17
18/// Rule to push down validity masking from MaskedArray parent into PrimitiveArray child.
19///
20/// When a PrimitiveArray is wrapped by a MaskedArray, this rule merges the mask's validity
21/// with the PrimitiveArray's existing validity, eliminating the need for the MaskedArray wrapper.
22#[derive(Default, Debug)]
23pub struct PrimitiveMaskedValidityRule;
24
25impl ArrayParentReduceRule<Exact<PrimitiveVTable>, Exact<MaskedVTable>>
26    for PrimitiveMaskedValidityRule
27{
28    fn child(&self) -> Exact<PrimitiveVTable> {
29        Exact::from(&PrimitiveVTable)
30    }
31
32    fn parent(&self) -> Exact<MaskedVTable> {
33        Exact::from(&MaskedVTable)
34    }
35
36    fn reduce_parent(
37        &self,
38        array: &PrimitiveArray,
39        parent: &MaskedArray,
40        _child_idx: usize,
41    ) -> VortexResult<Option<ArrayRef>> {
42        // Merge the parent's validity mask into the child's validity
43        // TODO(joe): make this lazy
44        let masked_array = match_each_native_ptype!(array.ptype(), |T| {
45            // SAFETY: Since we are only flipping some bits in the validity, all invariants that
46            // were upheld are still upheld.
47            unsafe {
48                PrimitiveArray::new_unchecked(
49                    Buffer::<T>::from_byte_buffer(array.byte_buffer().clone()),
50                    array.validity().clone().and(parent.validity().clone()),
51                )
52            }
53            .into_array()
54        });
55
56        Ok(Some(masked_array))
57    }
58}