vortex_runend/compute/
slice.rs

1use vortex_array::arrays::ConstantArray;
2use vortex_array::compute::{SliceFn, scalar_at, slice};
3use vortex_array::{Array, ArrayRef};
4use vortex_error::VortexResult;
5
6use crate::{RunEndArray, RunEndEncoding};
7
8impl SliceFn<&RunEndArray> for RunEndEncoding {
9    fn slice(&self, array: &RunEndArray, start: usize, stop: usize) -> VortexResult<ArrayRef> {
10        let new_length = stop - start;
11
12        let (slice_begin, slice_end) = if new_length == 0 {
13            let values_len = array.values().len();
14            (values_len, values_len)
15        } else {
16            let physical_start = array.find_physical_index(start)?;
17            let physical_stop = array.find_physical_index(stop)?;
18
19            (physical_start, physical_stop + 1)
20        };
21
22        if slice_begin + 1 == slice_end {
23            let value = scalar_at(array.values(), slice_begin)?;
24            return Ok(ConstantArray::new(value, new_length).into_array());
25        }
26
27        Ok(RunEndArray::with_offset_and_length(
28            slice(array.ends(), slice_begin, slice_end)?,
29            slice(array.values(), slice_begin, slice_end)?,
30            if new_length == 0 {
31                0
32            } else {
33                start + array.offset()
34            },
35            new_length,
36        )?
37        .into_array())
38    }
39}
40
41#[cfg(test)]
42mod tests {
43    use vortex_array::compute::{SliceFn, slice};
44    use vortex_array::{Array, ArrayStatistics, IntoArray, ToCanonical};
45    use vortex_buffer::buffer;
46    use vortex_dtype::{DType, Nullability, PType};
47
48    use crate::{RunEndArray, RunEndEncoding};
49
50    #[test]
51    fn slice_array() {
52        let arr = slice(
53            &RunEndArray::try_new(
54                buffer![2u32, 5, 10].into_array(),
55                buffer![1i32, 2, 3].into_array(),
56            )
57            .unwrap(),
58            3,
59            8,
60        )
61        .unwrap();
62        assert_eq!(
63            arr.dtype(),
64            &DType::Primitive(PType::I32, Nullability::NonNullable)
65        );
66        assert_eq!(arr.len(), 5);
67
68        assert_eq!(
69            arr.to_primitive().unwrap().as_slice::<i32>(),
70            vec![2, 2, 3, 3, 3]
71        );
72    }
73
74    #[test]
75    fn double_slice() {
76        let arr = slice(
77            &RunEndArray::try_new(
78                buffer![2u32, 5, 10].into_array(),
79                buffer![1i32, 2, 3].into_array(),
80            )
81            .unwrap(),
82            3,
83            8,
84        )
85        .unwrap();
86        assert_eq!(arr.len(), 5);
87
88        let doubly_sliced = slice(&arr, 0, 3).unwrap();
89
90        assert_eq!(
91            doubly_sliced.to_primitive().unwrap().as_slice::<i32>(),
92            vec![2, 2, 3]
93        );
94    }
95
96    #[test]
97    fn slice_end_inclusive() {
98        let arr = slice(
99            &RunEndArray::try_new(
100                buffer![2u32, 5, 10].into_array(),
101                buffer![1i32, 2, 3].into_array(),
102            )
103            .unwrap(),
104            4,
105            10,
106        )
107        .unwrap();
108        assert_eq!(
109            arr.dtype(),
110            &DType::Primitive(PType::I32, Nullability::NonNullable)
111        );
112        assert_eq!(arr.len(), 6);
113
114        assert_eq!(
115            arr.to_primitive().unwrap().as_slice::<i32>(),
116            vec![2, 3, 3, 3, 3, 3]
117        );
118    }
119
120    #[test]
121    fn slice_at_end() {
122        let re_array = RunEndArray::try_new(
123            buffer![7_u64, 10].into_array(),
124            buffer![2_u64, 3].into_array(),
125        )
126        .unwrap();
127
128        assert_eq!(re_array.len(), 10);
129
130        let sliced_array = RunEndEncoding
131            .slice(&re_array, re_array.len(), re_array.len())
132            .unwrap();
133        assert!(sliced_array.is_empty());
134
135        let re_slice = RunEndArray::try_from(sliced_array).unwrap();
136        assert!(re_slice.ends().is_empty());
137        assert!(re_slice.values().is_empty())
138    }
139
140    #[test]
141    fn slice_single_end() {
142        let re_array = RunEndArray::try_new(
143            buffer![7_u64, 10].into_array(),
144            buffer![2_u64, 3].into_array(),
145        )
146        .unwrap();
147
148        assert_eq!(re_array.len(), 10);
149
150        let sliced_array = RunEndEncoding.slice(&re_array, 2, 5).unwrap();
151
152        assert!(sliced_array.is_constant())
153    }
154}
vortex_runend/compute/slice.rs

vortex_runend/compute/
slice.rs
