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

vortex_alp/alp_rd/
ops.rs
