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

vortex_array/arrays/primitive/compute/
mask.rs
