vortex-fastlanes 0.83.0

Vortex fastlanes arrays
Documentation
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use vortex_array::ArrayView;
use vortex_array::ExecutionCtx;
use vortex_array::match_each_integer_ptype;
use vortex_array::scalar::Scalar;
use vortex_array::vtable::OperationsVTable;
use vortex_error::VortexExpect;
use vortex_error::VortexResult;

use super::FoR;
use crate::r#for::array::FoRArrayExt;
use crate::r#for::array::FoRArraySlotsExt;
impl OperationsVTable<FoR> for FoR {
    fn scalar_at(
        array: ArrayView<'_, FoR>,
        index: usize,
        ctx: &mut ExecutionCtx,
    ) -> VortexResult<Scalar> {
        let encoded_pvalue = array.encoded().execute_scalar(index, ctx)?;
        let encoded_pvalue = encoded_pvalue.as_primitive();
        let reference = array.reference_scalar();
        let reference = reference.as_primitive();

        Ok(match_each_integer_ptype!(array.ptype(), |P| {
            encoded_pvalue
                .typed_value::<P>()
                .map(|v| {
                    v.wrapping_add(
                        reference
                            .typed_value::<P>()
                            .vortex_expect("FoRArray Reference value cannot be null"),
                    )
                })
                .map(|v| Scalar::primitive::<P>(v, array.reference_scalar().dtype().nullability()))
                .unwrap_or_else(|| Scalar::null(array.reference_scalar().dtype().clone()))
        }))
    }
}

#[cfg(test)]
mod test {
    use std::sync::LazyLock;

    use vortex_array::VortexSessionExecute;
    use vortex_array::arrays::PrimitiveArray;
    use vortex_array::assert_arrays_eq;
    use vortex_session::VortexSession;

    use crate::FoRData;

    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
        let session = vortex_array::array_session();
        crate::initialize(&session);
        session
    });

    #[test]
    fn for_scalar_at() {
        let mut ctx = SESSION.create_execution_ctx();
        let for_arr = FoRData::encode(
            PrimitiveArray::from_iter([-100, 1100, 1500, 1900]),
            &mut ctx,
        )
        .unwrap();
        let expected = PrimitiveArray::from_iter([-100, 1100, 1500, 1900]);
        assert_arrays_eq!(for_arr, expected, &mut ctx);
    }
}