vortex_array/arrays/primitive/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;
6
7use crate::arrays::PrimitiveVTable;
8use crate::arrays::primitive::PrimitiveArray;
9use crate::compute::{MaskKernel, MaskKernelAdapter};
10use crate::vtable::ValidityHelper;
11use crate::{ArrayRef, IntoArray, register_kernel};
12
13impl MaskKernel for PrimitiveVTable {
14    fn mask(&self, array: &PrimitiveArray, mask: &Mask) -> VortexResult<ArrayRef> {
15        let validity = array.validity().mask(mask)?;
16        Ok(
17            PrimitiveArray::from_byte_buffer(array.byte_buffer().clone(), array.ptype(), validity)
18                .into_array(),
19        )
20    }
21}
22
23register_kernel!(MaskKernelAdapter(PrimitiveVTable).lift());
24
25#[cfg(test)]
26mod test {
27    use rstest::rstest;
28
29    use crate::arrays::PrimitiveArray;
30    use crate::compute::conformance::mask::test_mask_conformance;
31
32    #[rstest]
33    #[case(PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]))]
34    #[case(PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4), None]))]
35    #[case(PrimitiveArray::from_iter([42u64]))]
36    #[case(PrimitiveArray::from_iter(0..100i32))]
37    #[case(PrimitiveArray::from_option_iter((0..100).map(|i| if i % 5 == 0 { None } else { Some(i as i64) })))]
38    #[case(PrimitiveArray::from_iter([0.1f32, 0.2, 0.3, 0.4, 0.5]))]
39    #[case(PrimitiveArray::from_option_iter([Some(1.1f64), None, Some(2.2), Some(3.3), None]))]
40    fn test_mask_primitive_conformance(#[case] array: PrimitiveArray) {
41        test_mask_conformance(array.as_ref());
42    }
43}
vortex_array/arrays/primitive/compute/mask.rs

vortex_array/arrays/primitive/compute/
mask.rs
