vortex_fastlanes/for/vtable/
operations.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::ops::Range;
5
6use vortex_array::ArrayRef;
7use vortex_array::IntoArray;
8use vortex_array::vtable::OperationsVTable;
9use vortex_dtype::match_each_integer_ptype;
10use vortex_error::VortexExpect;
11use vortex_scalar::Scalar;
12
13use super::FoRVTable;
14use crate::FoRArray;
15
16impl OperationsVTable<FoRVTable> for FoRVTable {
17    fn slice(array: &FoRArray, range: Range<usize>) -> ArrayRef {
18        // SAFETY: Just slicing encoded data does not affect FOR.
19        unsafe {
20            FoRArray::new_unchecked(
21                array.encoded().slice(range),
22                array.reference_scalar().clone(),
23            )
24            .into_array()
25        }
26    }
27
28    fn scalar_at(array: &FoRArray, index: usize) -> Scalar {
29        let encoded_pvalue = array.encoded().scalar_at(index);
30        let encoded_pvalue = encoded_pvalue.as_primitive();
31        let reference = array.reference_scalar();
32        let reference = reference.as_primitive();
33
34        match_each_integer_ptype!(array.ptype(), |P| {
35            encoded_pvalue
36                .typed_value::<P>()
37                .map(|v| {
38                    v.wrapping_add(
39                        reference
40                            .typed_value::<P>()
41                            .vortex_expect("FoRArray Reference value cannot be null"),
42                    )
43                })
44                .map(|v| Scalar::primitive::<P>(v, array.reference_scalar().dtype().nullability()))
45                .unwrap_or_else(|| Scalar::null(array.reference_scalar().dtype().clone()))
46        })
47    }
48}
49
50#[cfg(test)]
51mod test {
52    use vortex_array::arrays::PrimitiveArray;
53    use vortex_array::assert_arrays_eq;
54
55    use crate::FoRArray;
56
57    #[test]
58    fn for_scalar_at() {
59        let for_arr =
60            FoRArray::encode(PrimitiveArray::from_iter([-100, 1100, 1500, 1900])).unwrap();
61        let expected = PrimitiveArray::from_iter([-100, 1100, 1500, 1900]);
62        assert_arrays_eq!(for_arr, expected);
63    }
64}
vortex_fastlanes/for/vtable/operations.rs

vortex_fastlanes/for/vtable/
operations.rs
