vortex_array/arrays/primitive/vtable/
rules.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::matchers::Exact;
15use crate::optimizer::rules::ArrayParentReduceRule;
16use crate::optimizer::rules::ParentRuleSet;
17use crate::vtable::ValidityHelper;
18
19pub(super) const RULES: ParentRuleSet<PrimitiveVTable> =
20    ParentRuleSet::new(&[ParentRuleSet::lift(&PrimitiveMaskedValidityRule)]);
21
22/// Rule to push down validity masking from MaskedArray parent into PrimitiveArray child.
23///
24/// When a PrimitiveArray is wrapped by a MaskedArray, this rule merges the mask's validity
25/// with the PrimitiveArray's existing validity, eliminating the need for the MaskedArray wrapper.
26#[derive(Default, Debug)]
27pub struct PrimitiveMaskedValidityRule;
28
29impl ArrayParentReduceRule<PrimitiveVTable> for PrimitiveMaskedValidityRule {
30    type Parent = Exact<MaskedVTable>;
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}
vortex_array/arrays/primitive/vtable/rules.rs

vortex_array/arrays/primitive/vtable/
rules.rs
