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::array::ArrayView;
9use crate::arrays::Primitive;
10use crate::arrays::PrimitiveArray;
11use crate::scalar_fn::fns::mask::MaskReduce;
12use crate::validity::Validity;
13
14impl MaskReduce for Primitive {
15    fn mask(array: ArrayView<'_, Primitive>, 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.validity()?.and(Validity::Array(mask.clone()))?,
22            )
23            .into_array()
24        }))
25    }
26}
27
28#[cfg(test)]
29mod test {
30    use rstest::rstest;
31
32    use crate::IntoArray;
33    use crate::arrays::PrimitiveArray;
34    use crate::compute::conformance::mask::test_mask_conformance;
35
36    #[rstest]
37    #[case(PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]))]
38    #[case(PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4), None]))]
39    #[case(PrimitiveArray::from_iter([42u64]))]
40    #[case(PrimitiveArray::from_iter(0..100i32))]
41    #[case(PrimitiveArray::from_option_iter((0..100).map(|i| if i % 5 == 0 { None } else { Some(i as i64) })))]
42    #[case(PrimitiveArray::from_iter([0.1f32, 0.2, 0.3, 0.4, 0.5]))]
43    #[case(PrimitiveArray::from_option_iter([Some(1.1f64), None, Some(2.2), Some(3.3), None]))]
44    fn test_mask_primitive_conformance(#[case] array: PrimitiveArray) {
45        test_mask_conformance(&array.into_array());
46    }
47}
vortex_array/arrays/primitive/compute/mask.rs

vortex_array/arrays/primitive/compute/
mask.rs
