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

vortex_array/arrays/masked/compute/
mask.rs
