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

vortex_alp/alp_rd/
ops.rs
