Skip to main content

vortex_array/arrays/dict/vtable/
operations.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexExpect;
5use vortex_error::VortexResult;
6use vortex_error::vortex_err;
7
8use super::Dict;
9use crate::ExecutionCtx;
10use crate::array::ArrayView;
11use crate::array::OperationsVTable;
12use crate::array::ProbeState;
13use crate::arrays::dict::DictSlots;
14use crate::scalar::Scalar;
15
16impl OperationsVTable<Dict> for Dict {
17    type ProbeState = ();
18
19    fn probe_scalar(
20        state: &mut ProbeState<'_, Dict>,
21        index: usize,
22        ctx: &mut ExecutionCtx,
23    ) -> VortexResult<Scalar> {
24        let dtype = state.array().dtype().clone();
25        let code = state
26            .slot(DictSlots::CODES)?
27            .ok_or_else(|| vortex_err!("Dict codes slot is missing"))?
28            .execute_scalar(index, ctx)?;
29        let Some(dict_index) = code.as_primitive().as_::<usize>() else {
30            return Ok(Scalar::null(dtype));
31        };
32
33        Ok(state
34            .slot(DictSlots::VALUES)?
35            .ok_or_else(|| vortex_err!("Dict values slot is missing"))?
36            .execute_scalar(dict_index, ctx)?
37            .cast(&dtype)
38            .vortex_expect("Array dtype will only differ by nullability"))
39    }
40
41    fn scalar_at(
42        array: ArrayView<'_, Dict>,
43        index: usize,
44        ctx: &mut ExecutionCtx,
45    ) -> VortexResult<Scalar> {
46        Self::probe_scalar(&mut ProbeState::once(array), index, ctx)
47    }
48}