vortex_dict/
ops.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::ops::Range;
5
6use vortex_array::arrays::{ConstantArray, ConstantVTable};
7use vortex_array::vtable::OperationsVTable;
8use vortex_array::{Array, ArrayRef, IntoArray};
9use vortex_error::VortexExpect;
10use vortex_scalar::Scalar;
11
12use crate::{DictArray, DictVTable};
13
14impl OperationsVTable<DictVTable> for DictVTable {
15    fn slice(array: &DictArray, range: Range<usize>) -> ArrayRef {
16        let sliced_code = array.codes().slice(range);
17        if sliced_code.is::<ConstantVTable>() {
18            let code = &sliced_code.scalar_at(0).as_primitive().as_::<usize>();
19            return if let Some(code) = code {
20                ConstantArray::new(array.values().scalar_at(*code), sliced_code.len()).into_array()
21            } else {
22                ConstantArray::new(Scalar::null(array.dtype().clone()), sliced_code.len())
23                    .to_array()
24            };
25        }
26        // SAFETY: slicing the codes preserves invariants
27        unsafe { DictArray::new_unchecked(sliced_code, array.values().clone()).into_array() }
28    }
29
30    fn scalar_at(array: &DictArray, index: usize) -> Scalar {
31        let Some(dict_index) = array.codes().scalar_at(index).as_primitive().as_::<usize>() else {
32            return Scalar::null(array.dtype().clone());
33        };
34
35        array
36            .values()
37            .scalar_at(dict_index)
38            .cast(array.dtype())
39            .vortex_expect("Array dtype will only differ by nullability")
40    }
41}
42
43#[cfg(test)]
44mod tests {
45    use vortex_array::IntoArray;
46    use vortex_array::arrays::PrimitiveArray;
47    use vortex_buffer::buffer;
48    use vortex_dtype::Nullability;
49    use vortex_scalar::Scalar;
50
51    use crate::DictArray;
52
53    #[test]
54    fn test_slice_into_const_dict() {
55        let dict = DictArray::try_new(
56            PrimitiveArray::from_option_iter(vec![Some(0u32), None, Some(1)]).to_array(),
57            PrimitiveArray::from_option_iter(vec![Some(0i32), Some(1), Some(2)]).to_array(),
58        )
59        .unwrap();
60
61        assert_eq!(
62            Some(Scalar::new(dict.dtype().clone(), 0i32.into())),
63            dict.slice(0..1).as_constant()
64        );
65
66        assert_eq!(
67            Some(Scalar::null(dict.dtype().clone())),
68            dict.slice(1..2).as_constant()
69        );
70    }
71
72    #[test]
73    fn test_scalar_at_null_code() {
74        let dict = DictArray::try_new(
75            PrimitiveArray::from_option_iter(vec![None, Some(0u32), None]).to_array(),
76            buffer![1i32].into_array(),
77        )
78        .unwrap();
79
80        assert_eq!(dict.scalar_at(0), Scalar::null(dict.dtype().clone()));
81        assert_eq!(
82            dict.scalar_at(1),
83            Scalar::primitive(1, Nullability::Nullable)
84        );
85    }
86}
vortex_dict/ops.rs

vortex_dict/
ops.rs
