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

vortex_array/arrays/primitive/compute/
mask.rs
