vortex_array/arrays/dict/vtable/
operations.rs1use 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}