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                let dtype = array.values().dtype().with_nullability(
23                    array.values().dtype().nullability() | array.codes().dtype().nullability(),
24                );
25                ConstantArray::new(Scalar::null(dtype), sliced_code.len()).to_array()
26            };
27        }
28        // SAFETY: slicing the codes preserves invariants
29        unsafe { DictArray::new_unchecked(sliced_code, array.values().clone()).into_array() }
30    }
31
32    fn scalar_at(array: &DictArray, index: usize) -> Scalar {
33        let dict_index: usize = array
34            .codes()
35            .scalar_at(index)
36            .as_ref()
37            .try_into()
38            .vortex_expect("code overflowed usize");
39        array.values().scalar_at(dict_index)
40    }
41}
42
43#[cfg(test)]
44mod tests {
45    use vortex_array::arrays::PrimitiveArray;
46    use vortex_scalar::Scalar;
47
48    use crate::DictArray;
49
50    #[test]
51    fn test_slice_into_const_dict() {
52        let dict = DictArray::try_new(
53            PrimitiveArray::from_option_iter(vec![Some(0u32), None, Some(1)]).to_array(),
54            PrimitiveArray::from_option_iter(vec![Some(0i32), Some(1), Some(2)]).to_array(),
55        )
56        .unwrap();
57
58        assert_eq!(
59            Some(Scalar::new(dict.dtype().clone(), 0i32.into())),
60            dict.slice(0..1).as_constant()
61        );
62
63        assert_eq!(
64            Some(Scalar::null(dict.dtype().clone())),
65            dict.slice(1..2).as_constant()
66        );
67    }
68}
vortex_dict/ops.rs

vortex_dict/
ops.rs
