vortex_array/arrays/dict/vtable/
operations.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::ops::Range;
5
6use vortex_error::VortexExpect;
7use vortex_scalar::Scalar;
8
9use super::DictVTable;
10use crate::Array;
11use crate::ArrayRef;
12use crate::IntoArray;
13use crate::arrays::ConstantArray;
14use crate::arrays::ConstantVTable;
15use crate::arrays::dict::DictArray;
16use crate::vtable::OperationsVTable;
17
18impl OperationsVTable<DictVTable> for DictVTable {
19    fn slice(array: &DictArray, range: Range<usize>) -> ArrayRef {
20        let sliced_code = array.codes().slice(range);
21        if sliced_code.is::<ConstantVTable>() {
22            let code = &sliced_code.scalar_at(0).as_primitive().as_::<usize>();
23            return if let Some(code) = code {
24                ConstantArray::new(array.values().scalar_at(*code), sliced_code.len()).into_array()
25            } else {
26                ConstantArray::new(Scalar::null(array.dtype().clone()), sliced_code.len())
27                    .to_array()
28            };
29        }
30        // SAFETY: slicing the codes preserves invariants.
31        unsafe { DictArray::new_unchecked(sliced_code, array.values().clone()).into_array() }
32    }
33
34    fn scalar_at(array: &DictArray, index: usize) -> Scalar {
35        let Some(dict_index) = array.codes().scalar_at(index).as_primitive().as_::<usize>() else {
36            return Scalar::null(array.dtype().clone());
37        };
38
39        array
40            .values()
41            .scalar_at(dict_index)
42            .cast(array.dtype())
43            .vortex_expect("Array dtype will only differ by nullability")
44    }
45}