Skip to main content

vortex_array/arrays/extension/vtable/
operations.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5use vortex_error::vortex_err;
6
7use crate::ExecutionCtx;
8use crate::array::ArrayView;
9use crate::array::OperationsVTable;
10use crate::array::ProbeState;
11use crate::arrays::Extension;
12use crate::arrays::extension::ExtensionArrayExt;
13use crate::arrays::extension::ExtensionSlots;
14use crate::scalar::Scalar;
15
16impl OperationsVTable<Extension> for Extension {
17    type ProbeState = ();
18
19    fn probe_scalar(
20        state: &mut ProbeState<'_, Extension>,
21        index: usize,
22        ctx: &mut ExecutionCtx,
23    ) -> VortexResult<Scalar> {
24        let ext_dtype = state.array().ext_dtype().clone();
25        let storage = state
26            .slot(ExtensionSlots::STORAGE)?
27            .ok_or_else(|| vortex_err!("Extension storage slot is missing"))?
28            .execute_scalar(index, ctx)?;
29        Ok(Scalar::extension_ref(ext_dtype, storage))
30    }
31
32    fn scalar_at(
33        array: ArrayView<'_, Extension>,
34        index: usize,
35        ctx: &mut ExecutionCtx,
36    ) -> VortexResult<Scalar> {
37        Self::probe_scalar(&mut ProbeState::once(array), index, ctx)
38    }
39}