vortex_runend/compute/
slice.rs

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

vortex_runend/compute/
slice.rs
