Skip to main content

vortex_fastlanes/for/vtable/
operations.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::match_each_integer_ptype;
5use vortex_array::scalar::Scalar;
6use vortex_array::vtable::OperationsVTable;
7use vortex_error::VortexExpect;
8use vortex_error::VortexResult;
9
10use super::FoRVTable;
11use crate::FoRArray;
12
13impl OperationsVTable<FoRVTable> for FoRVTable {
14    fn scalar_at(array: &FoRArray, index: usize) -> VortexResult<Scalar> {
15        let encoded_pvalue = array.encoded().scalar_at(index)?;
16        let encoded_pvalue = encoded_pvalue.as_primitive();
17        let reference = array.reference_scalar();
18        let reference = reference.as_primitive();
19
20        Ok(match_each_integer_ptype!(array.ptype(), |P| {
21            encoded_pvalue
22                .typed_value::<P>()
23                .map(|v| {
24                    v.wrapping_add(
25                        reference
26                            .typed_value::<P>()
27                            .vortex_expect("FoRArray Reference value cannot be null"),
28                    )
29                })
30                .map(|v| Scalar::primitive::<P>(v, array.reference_scalar().dtype().nullability()))
31                .unwrap_or_else(|| Scalar::null(array.reference_scalar().dtype().clone()))
32        }))
33    }
34}
35
36#[cfg(test)]
37mod test {
38    use vortex_array::arrays::PrimitiveArray;
39    use vortex_array::assert_arrays_eq;
40
41    use crate::FoRArray;
42
43    #[test]
44    fn for_scalar_at() {
45        let for_arr =
46            FoRArray::encode(PrimitiveArray::from_iter([-100, 1100, 1500, 1900])).unwrap();
47        let expected = PrimitiveArray::from_iter([-100, 1100, 1500, 1900]);
48        assert_arrays_eq!(for_arr, expected);
49    }
50}