Skip to main content

vortex_alp/alp_rd/
ops.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::Array;
5use vortex_array::scalar::Scalar;
6use vortex_array::vtable::OperationsVTable;
7use vortex_error::VortexExpect;
8use vortex_error::VortexResult;
9
10use crate::ALPRDArray;
11use crate::ALPRDVTable;
12
13impl OperationsVTable<ALPRDVTable> for ALPRDVTable {
14    fn scalar_at(array: &ALPRDArray, index: usize) -> VortexResult<Scalar> {
15        // The left value can either be a direct value, or an exception.
16        // The exceptions array represents exception positions with non-null values.
17        let maybe_patched_value = match array.left_parts_patches() {
18            Some(patches) => patches.get_patched(index)?,
19            None => None,
20        };
21        let left = match maybe_patched_value {
22            Some(patched_value) => patched_value
23                .as_primitive()
24                .as_::<u16>()
25                .vortex_expect("patched values must be non-null"),
26            _ => {
27                let left_code: u16 = array
28                    .left_parts()
29                    .scalar_at(index)?
30                    .as_primitive()
31                    .as_::<u16>()
32                    .vortex_expect("left_code must be non-null");
33                array.left_parts_dictionary()[left_code as usize]
34            }
35        };
36
37        // combine left and right values
38        Ok(if array.is_f32() {
39            let right: u32 = array
40                .right_parts()
41                .scalar_at(index)?
42                .as_primitive()
43                .as_::<u32>()
44                .vortex_expect("non-null");
45            let packed = f32::from_bits((left as u32) << array.right_bit_width() | right);
46            Scalar::primitive(packed, array.dtype().nullability())
47        } else {
48            let right: u64 = array
49                .right_parts()
50                .scalar_at(index)?
51                .as_primitive()
52                .as_::<u64>()
53                .vortex_expect("non-null");
54            let packed = f64::from_bits(((left as u64) << array.right_bit_width()) | right);
55            Scalar::primitive(packed, array.dtype().nullability())
56        })
57    }
58}
59
60#[cfg(test)]
61mod test {
62    use rstest::rstest;
63    use vortex_array::arrays::PrimitiveArray;
64    use vortex_array::assert_arrays_eq;
65    use vortex_array::scalar::Scalar;
66
67    use crate::ALPRDFloat;
68    use crate::RDEncoder;
69
70    #[rstest]
71    #[case(0.1f32, 0.2f32, 3e25f32)]
72    #[case(0.1f64, 0.2f64, 3e100f64)]
73    fn test_slice<T: ALPRDFloat>(#[case] a: T, #[case] b: T, #[case] outlier: T) {
74        let array = PrimitiveArray::from_iter([a, b, outlier]);
75        let encoded = RDEncoder::new(&[a, b]).encode(&array);
76
77        assert!(encoded.left_parts_patches().is_some());
78        assert_arrays_eq!(encoded, array);
79    }
80
81    #[rstest]
82    #[case(0.1f32, 0.2f32, 3e25f32)]
83    #[case(0.1f64, 0.2f64, 3e100f64)]
84    fn test_scalar_at<T: ALPRDFloat + Into<Scalar>>(
85        #[case] a: T,
86        #[case] b: T,
87        #[case] outlier: T,
88    ) {
89        let array = PrimitiveArray::from_iter([a, b, outlier]);
90        let encoded = RDEncoder::new(&[a, b]).encode(&array);
91        assert!(encoded.left_parts_patches().is_some());
92        assert_arrays_eq!(encoded, array);
93    }
94
95    #[test]
96    fn nullable_scalar_at() {
97        let a = 0.1f64;
98        let b = 0.2f64;
99        let outlier = 3e100f64;
100        let array = PrimitiveArray::from_option_iter([Some(a), Some(b), Some(outlier)]);
101        let encoded = RDEncoder::new(&[a, b]).encode(&array);
102        assert!(encoded.left_parts_patches().is_some());
103        assert_arrays_eq!(
104            encoded,
105            PrimitiveArray::from_option_iter([Some(a), Some(b), Some(outlier)])
106        );
107    }
108}