vortex_alp/alp/compute/
mask.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::ArrayRef;
5use vortex_array::ArrayView;
6use vortex_array::ExecutionCtx;
7use vortex_array::IntoArray;
8use vortex_array::builtins::ArrayBuiltins;
9use vortex_array::scalar_fn::fns::mask::MaskKernel;
10use vortex_array::scalar_fn::fns::mask::MaskReduce;
11use vortex_array::validity::Validity;
12use vortex_error::VortexResult;
13
14use crate::ALP;
15use crate::ALPArrayExt;
16use crate::ALPArraySlotsExt;
17
18impl MaskReduce for ALP {
19    fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
20        // Masking sparse patches requires reading indices, fall back to kernel.
21        if array.patches().is_some() {
22            return Ok(None);
23        }
24        let masked_encoded = array.encoded().clone().mask(mask.clone())?;
25        Ok(Some(
26            ALP::new(masked_encoded, array.exponents(), None).into_array(),
27        ))
28    }
29}
30
31impl MaskKernel for ALP {
32    fn mask(
33        array: ArrayView<'_, Self>,
34        mask: &ArrayRef,
35        ctx: &mut ExecutionCtx,
36    ) -> VortexResult<Option<ArrayRef>> {
37        let vortex_mask = Validity::Array(mask.not()?).execute_mask(array.len(), ctx)?;
38        let masked_encoded = array.encoded().clone().mask(mask.clone())?;
39        let masked_patches = array
40            .patches()
41            .map(|p| p.mask(&vortex_mask, ctx))
42            .transpose()?
43            .flatten();
44        Ok(Some(
45            ALP::try_new(masked_encoded, array.exponents(), masked_patches)?.into_array(),
46        ))
47    }
48}
49
50#[cfg(test)]
51mod test {
52    use rstest::rstest;
53    use vortex_array::IntoArray;
54    use vortex_array::VortexSessionExecute;
55    use vortex_array::array_session;
56    use vortex_array::arrays::BoolArray;
57    use vortex_array::arrays::PrimitiveArray;
58    use vortex_array::compute::conformance::mask::test_mask_conformance;
59    use vortex_array::dtype::Nullability;
60    use vortex_array::scalar_fn::fns::mask::MaskKernel;
61    use vortex_buffer::buffer;
62
63    use crate::alp::array::ALPArrayExt;
64    use crate::alp_encode;
65
66    #[rstest]
67    #[case(buffer![10.5f32, 20.5, 30.5, 40.5, 50.5].into_array())]
68    #[case(buffer![1000.123f64, 2000.456, 3000.789, 4000.012, 5000.345].into_array())]
69    #[case(PrimitiveArray::from_option_iter([Some(1.1f32), None, Some(2.2), Some(3.3), None]).into_array())]
70    #[case(buffer![99.99f64].into_array())]
71    #[case(buffer![
72        0.1f32, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0,
73        1.1, 1.2, 1.3, 1.4, 1.5, 1.6, 1.7, 1.8, 1.9, 2.0
74    ].into_array())]
75    fn test_mask_alp_conformance(#[case] array: vortex_array::ArrayRef) {
76        let mut ctx = array_session().create_execution_ctx();
77        let array_primitive = array.execute::<PrimitiveArray>(&mut ctx).unwrap();
78        let alp = alp_encode(array_primitive.as_view(), None, &mut ctx).unwrap();
79        test_mask_conformance(&alp.into_array(), &mut ctx);
80    }
81
82    #[test]
83    fn test_mask_alp_with_patches() {
84        use std::f64::consts::PI;
85        let mut ctx = array_session().create_execution_ctx();
86        // PI doesn't encode cleanly with ALP, so it creates patches.
87        let values: Vec<f64> = (0..100)
88            .map(|i| if i % 4 == 3 { PI } else { 1.0 })
89            .collect();
90        let array = PrimitiveArray::from_iter(values);
91        let alp = alp_encode(array.as_view(), None, &mut ctx).unwrap();
92        assert!(alp.patches().is_some(), "expected patches");
93        test_mask_conformance(&alp.into_array(), &mut ctx);
94    }
95
96    #[test]
97    fn test_mask_alp_with_patches_casts_surviving_patch_values_to_nullable() {
98        let mut ctx = array_session().create_execution_ctx();
99        let values = PrimitiveArray::from_iter([1.234f32, f32::NAN, 2.345, f32::INFINITY, 3.456]);
100        let alp = alp_encode(values.as_view(), None, &mut ctx).unwrap();
101        assert!(alp.patches().is_some(), "expected patches");
102
103        let keep_mask = BoolArray::from_iter([false, true, true, true, true]).into_array();
104        let masked = <crate::ALP as MaskKernel>::mask(alp.as_view(), &keep_mask, &mut ctx)
105            .unwrap()
106            .unwrap();
107
108        let masked_alp = masked.as_opt::<crate::ALP>().unwrap();
109        let masked_patches = masked_alp.patches().unwrap();
110
111        assert_eq!(masked.dtype().nullability(), Nullability::Nullable);
112        assert_eq!(masked_patches.dtype().nullability(), Nullability::Nullable);
113    }
114}
vortex_alp/alp/compute/mask.rs

vortex_alp/alp/compute/
mask.rs
