vortex_dict/
ops.rs

1use vortex_array::arrays::{ConstantArray, ConstantVTable};
2use vortex_array::vtable::OperationsVTable;
3use vortex_array::{Array, ArrayRef, IntoArray};
4use vortex_error::VortexResult;
5use vortex_scalar::Scalar;
6
7use crate::{DictArray, DictVTable};
8
9impl OperationsVTable<DictVTable> for DictVTable {
10    fn slice(array: &DictArray, start: usize, stop: usize) -> VortexResult<ArrayRef> {
11        let sliced_code = array.codes().slice(start, stop)?;
12        if sliced_code.is::<ConstantVTable>() {
13            let code = Option::<usize>::try_from(&sliced_code.scalar_at(0)?)?;
14            return if let Some(code) = code {
15                Ok(
16                    ConstantArray::new(array.values().scalar_at(code)?, sliced_code.len())
17                        .to_array(),
18                )
19            } else {
20                let dtype = array.values().dtype().with_nullability(
21                    array.values().dtype().nullability() | array.codes().dtype().nullability(),
22                );
23                Ok(ConstantArray::new(Scalar::null(dtype), sliced_code.len()).to_array())
24            };
25        }
26        DictArray::try_new(sliced_code, array.values().clone()).map(|a| a.into_array())
27    }
28
29    fn scalar_at(array: &DictArray, index: usize) -> VortexResult<Scalar> {
30        let dict_index: usize = array.codes().scalar_at(index)?.as_ref().try_into()?;
31        array.values().scalar_at(dict_index)
32    }
33}
34
35#[cfg(test)]
36mod tests {
37    use vortex_array::arrays::PrimitiveArray;
38    use vortex_scalar::Scalar;
39
40    use crate::DictArray;
41
42    #[test]
43    fn test_slice_into_const_dict() {
44        let dict = DictArray::try_new(
45            PrimitiveArray::from_option_iter(vec![Some(0u32), None, Some(1)]).to_array(),
46            PrimitiveArray::from_option_iter(vec![Some(0i32), Some(1), Some(2)]).to_array(),
47        )
48        .unwrap();
49
50        assert_eq!(
51            Some(Scalar::new(dict.dtype().clone(), 0i32.into())),
52            dict.slice(0, 1).unwrap().as_constant()
53        );
54
55        assert_eq!(
56            Some(Scalar::null(dict.dtype().clone())),
57            dict.slice(1, 2).unwrap().as_constant()
58        );
59    }
60}
vortex_dict/ops.rs

vortex_dict/
ops.rs
