vortex_array/arrays/dict/compute/
slice.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_error::VortexResult;
8
9use crate::ArrayRef;
10use crate::IntoArray;
11use crate::array::ArrayView;
12use crate::arrays::Constant;
13use crate::arrays::ConstantArray;
14use crate::arrays::Dict;
15use crate::arrays::DictArray;
16use crate::arrays::Primitive;
17use crate::arrays::dict::DictArraySlotsExt;
18use crate::arrays::slice::SliceReduce;
19use crate::expr::stats::Precision;
20use crate::expr::stats::Stat;
21use crate::scalar::Scalar;
22use crate::scalar::ScalarValue;
23
24impl SliceReduce for Dict {
25    fn slice(array: ArrayView<'_, Self>, range: Range<usize>) -> VortexResult<Option<ArrayRef>> {
26        if let Some(code) = array.codes().as_opt::<Constant>() {
27            return slice_constant_code(array, code.scalar(), range.len());
28        }
29
30        let sliced_code = if let Some(codes) = array.codes().as_typed::<Primitive>() {
31            let sliced_code = <Primitive as SliceReduce>::slice(codes, range)?
32                .vortex_expect("Primitive SliceReduce should always return Some");
33            // Because we specialize the primitive branch here, we have to make sure to handle the stat inheritance
34            inherit_slice_stats(array.codes(), &sliced_code);
35            sliced_code
36        } else {
37            array.codes().slice(range)?
38        };
39
40        // TODO(joe): if the range is size 1 replace with a constant array
41        if let Some(code) = sliced_code.as_opt::<Constant>() {
42            return slice_constant_code(array, code.scalar(), sliced_code.len());
43        }
44        // SAFETY: slicing the codes preserves invariants.
45        let array =
46            unsafe { DictArray::new_unchecked(sliced_code, array.values().clone()).into_array() };
47
48        Ok(Some(array))
49    }
50}
51
52fn inherit_slice_stats(source: &ArrayRef, sliced: &ArrayRef) {
53    source.statistics().with_iter(|iter| {
54        sliced
55            .statistics()
56            .inherit(iter.filter(|(stat, value)| is_inheritable_true_slice_stat(*stat, value)));
57    });
58}
59
60fn is_inheritable_true_slice_stat(stat: Stat, value: &Precision<ScalarValue>) -> bool {
61    matches!(
62        stat,
63        Stat::IsConstant | Stat::IsSorted | Stat::IsStrictSorted
64    ) && value
65        .as_ref()
66        .as_exact()
67        .is_some_and(|value| matches!(value, ScalarValue::Bool(true)))
68}
69
70fn slice_constant_code(
71    array: ArrayView<'_, Dict>,
72    code: &Scalar,
73    len: usize,
74) -> VortexResult<Option<ArrayRef>> {
75    let code = code.as_primitive().as_::<usize>();
76    if let Some(code) = code {
77        let values = array.values().slice(code..code + 1)?;
78        Ok(Some(
79            DictArray::new(ConstantArray::new(0u8, len).into_array(), values).into_array(),
80        ))
81    } else {
82        Ok(Some(
83            ConstantArray::new(Scalar::null(array.dtype().clone()), len).into_array(),
84        ))
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use vortex_buffer::buffer;
91    use vortex_error::VortexResult;
92
93    use crate::IntoArray;
94    use crate::VortexSessionExecute;
95    use crate::array_session;
96    use crate::arrays::DictArray;
97    use crate::arrays::PrimitiveArray;
98    use crate::arrays::dict::compute::slice::ConstantArray;
99    use crate::assert_arrays_eq;
100    use crate::dtype::DType;
101    use crate::dtype::Nullability::Nullable;
102    use crate::dtype::PType;
103    use crate::scalar::Scalar;
104
105    #[test]
106    fn slice_constant_valid_code() -> VortexResult<()> {
107        let mut ctx = array_session().create_execution_ctx();
108        let dict = DictArray::new(
109            ConstantArray::new(1u8, 5).into_array(),
110            buffer![10i32, 20, 30].into_array(),
111        );
112        let sliced = dict.slice(1..4)?;
113        let expected = PrimitiveArray::from_iter([20i32, 20, 20]).into_array();
114        assert_arrays_eq!(sliced, expected, &mut ctx);
115        Ok(())
116    }
117
118    #[test]
119    fn slice_constant_null_code() -> VortexResult<()> {
120        let mut ctx = array_session().create_execution_ctx();
121        let dict = DictArray::new(
122            ConstantArray::new(Scalar::null(DType::Primitive(PType::U8, Nullable)), 5).into_array(),
123            buffer![10i32, 20, 30].into_array(),
124        );
125        let sliced = dict.slice(1..4)?;
126        let expected =
127            PrimitiveArray::from_option_iter([Option::<i32>::None, None, None]).into_array();
128        assert_arrays_eq!(sliced, expected, &mut ctx);
129        Ok(())
130    }
131}
vortex_array/arrays/dict/compute/slice.rs

vortex_array/arrays/dict/compute/
slice.rs
