vortex_array/arrays/masked/compute/
mask.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5use vortex_mask::Mask as MaskType;
6
7use crate::arrays::{MaskedArray, MaskedVTable};
8use crate::compute::{MaskKernel, MaskKernelAdapter};
9use crate::vtable::ValidityHelper;
10use crate::{ArrayRef, IntoArray, register_kernel};
11
12impl MaskKernel for MaskedVTable {
13    fn mask(&self, array: &MaskedArray, mask_arg: &MaskType) -> VortexResult<ArrayRef> {
14        // Combine the mask with the existing validity
15        // The child remains unchanged (no nulls), only validity is updated
16        let combined_validity = array.validity().mask(mask_arg);
17
18        Ok(MaskedArray::try_new(array.child.clone(), combined_validity)?.into_array())
19    }
20}
21
22register_kernel!(MaskKernelAdapter(MaskedVTable).lift());
23
24#[cfg(test)]
25mod tests {
26    use rstest::rstest;
27
28    use crate::IntoArray;
29    use crate::arrays::{MaskedArray, PrimitiveArray};
30    use crate::compute::conformance::mask::test_mask_conformance;
31    use crate::validity::Validity;
32
33    #[rstest]
34    #[case(
35        MaskedArray::try_new(
36            PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]).into_array(),
37            Validity::from_iter([true, true, false, true, false])
38        ).unwrap()
39    )]
40    #[case(
41        MaskedArray::try_new(
42            PrimitiveArray::from_iter([10i32, 20, 30]).into_array(),
43            Validity::AllValid
44        ).unwrap()
45    )]
46    #[case(
47        MaskedArray::try_new(
48            PrimitiveArray::from_iter(0..100).into_array(),
49            Validity::from_iter((0..100).map(|i| i % 3 != 0))
50        ).unwrap()
51    )]
52    fn test_mask_masked_conformance(#[case] array: MaskedArray) {
53        test_mask_conformance(array.as_ref());
54    }
55}
vortex_array/arrays/masked/compute/mask.rs

vortex_array/arrays/masked/compute/
mask.rs
