vortex_alp/alp_rd/
ops.rs

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

vortex_alp/alp_rd/
ops.rs
