vortex_runend/
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::Array;
7use vortex_array::ArrayRef;
8use vortex_array::IntoArray;
9use vortex_array::arrays::ConstantArray;
10use vortex_array::search_sorted::SearchResult;
11use vortex_array::search_sorted::SearchSorted;
12use vortex_array::search_sorted::SearchSortedSide;
13use vortex_array::vtable::OperationsVTable;
14use vortex_scalar::PValue;
15use vortex_scalar::Scalar;
16
17use crate::RunEndArray;
18use crate::RunEndVTable;
19
20impl OperationsVTable<RunEndVTable> for RunEndVTable {
21    fn slice(array: &RunEndArray, range: Range<usize>) -> ArrayRef {
22        let new_length = range.len();
23
24        let slice_begin = array.find_physical_index(range.start);
25        let slice_end = find_slice_end_index(array.ends(), range.end + array.offset());
26
27        // If the sliced range contains only a single run, opt to return a ConstantArray.
28        if slice_begin + 1 == slice_end {
29            let value = array.values().scalar_at(slice_begin);
30            return ConstantArray::new(value, new_length).into_array();
31        }
32
33        // SAFETY: we maintain the ends invariant in our slice implementation
34        unsafe {
35            RunEndArray::new_unchecked(
36                array.ends().slice(slice_begin..slice_end),
37                array.values().slice(slice_begin..slice_end),
38                range.start + array.offset(),
39                new_length,
40            )
41            .into_array()
42        }
43    }
44
45    fn scalar_at(array: &RunEndArray, index: usize) -> Scalar {
46        array.values().scalar_at(array.find_physical_index(index))
47    }
48}
49
50/// Find the physical offset for and index that would be an end of the slice i.e., one past the last element.
51///
52/// If the index exists in the array we want to take that position (as we are searching from the right)
53/// otherwise we want to take the next one
54pub(crate) fn find_slice_end_index(array: &dyn Array, index: usize) -> usize {
55    let result = array
56        .as_primitive_typed()
57        .search_sorted(&PValue::from(index), SearchSortedSide::Right);
58    match result {
59        SearchResult::Found(i) => i,
60        SearchResult::NotFound(i) => {
61            if i == array.len() {
62                i
63            } else {
64                i + 1
65            }
66        }
67    }
68}
69
70#[cfg(test)]
71mod tests {
72
73    use vortex_array::Array;
74    use vortex_array::IntoArray;
75    use vortex_array::arrays::PrimitiveArray;
76    use vortex_array::assert_arrays_eq;
77    use vortex_buffer::buffer;
78    use vortex_dtype::DType;
79    use vortex_dtype::Nullability;
80    use vortex_dtype::PType;
81
82    use crate::RunEndArray;
83
84    #[test]
85    fn slice_array() {
86        let arr = RunEndArray::try_new(
87            buffer![2u32, 5, 10].into_array(),
88            buffer![1i32, 2, 3].into_array(),
89        )
90        .unwrap()
91        .slice(3..8);
92        assert_eq!(
93            arr.dtype(),
94            &DType::Primitive(PType::I32, Nullability::NonNullable)
95        );
96        assert_eq!(arr.len(), 5);
97
98        let expected = PrimitiveArray::from_iter(vec![2i32, 2, 3, 3, 3]).into_array();
99        assert_arrays_eq!(arr, expected);
100    }
101
102    #[test]
103    fn double_slice() {
104        let arr = RunEndArray::try_new(
105            buffer![2u32, 5, 10].into_array(),
106            buffer![1i32, 2, 3].into_array(),
107        )
108        .unwrap()
109        .slice(3..8);
110        assert_eq!(arr.len(), 5);
111
112        let doubly_sliced = arr.slice(0..3);
113
114        let expected = PrimitiveArray::from_iter(vec![2i32, 2, 3]).into_array();
115        assert_arrays_eq!(doubly_sliced, expected);
116    }
117
118    #[test]
119    fn slice_end_inclusive() {
120        let arr = RunEndArray::try_new(
121            buffer![2u32, 5, 10].into_array(),
122            buffer![1i32, 2, 3].into_array(),
123        )
124        .unwrap()
125        .slice(4..10);
126        assert_eq!(
127            arr.dtype(),
128            &DType::Primitive(PType::I32, Nullability::NonNullable)
129        );
130        assert_eq!(arr.len(), 6);
131
132        let expected = PrimitiveArray::from_iter(vec![2i32, 3, 3, 3, 3, 3]).into_array();
133        assert_arrays_eq!(arr, expected);
134    }
135
136    #[test]
137    fn slice_at_end() {
138        let re_array = RunEndArray::try_new(
139            buffer![7_u64, 10].into_array(),
140            buffer![2_u64, 3].into_array(),
141        )
142        .unwrap();
143
144        assert_eq!(re_array.len(), 10);
145
146        let sliced_array = re_array.slice(re_array.len()..re_array.len());
147        assert!(sliced_array.is_empty());
148    }
149
150    #[test]
151    fn slice_single_end() {
152        let re_array = RunEndArray::try_new(
153            buffer![7_u64, 10].into_array(),
154            buffer![2_u64, 3].into_array(),
155        )
156        .unwrap();
157
158        assert_eq!(re_array.len(), 10);
159
160        let sliced_array = re_array.slice(2..5);
161
162        assert!(sliced_array.is_constant())
163    }
164
165    #[test]
166    fn ree_scalar_at_end() {
167        let scalar = RunEndArray::encode(buffer![1, 1, 1, 4, 4, 4, 2, 2, 5, 5, 5, 5].into_array())
168            .unwrap()
169            .scalar_at(11);
170        assert_eq!(scalar, 5.into());
171    }
172
173    #[test]
174    #[allow(clippy::cognitive_complexity)]
175    fn slice_along_run_boundaries() {
176        // Create a runend array with runs: [1, 1, 1] [4, 4, 4] [2, 2] [5, 5, 5, 5]
177        // Run ends at indices: 3, 6, 8, 12
178        let arr = RunEndArray::try_new(
179            buffer![3u32, 6, 8, 12].into_array(),
180            buffer![1i32, 4, 2, 5].into_array(),
181        )
182        .unwrap();
183
184        // Slice from start of first run to end of first run (indices 0..3)
185        let slice1 = arr.slice(0..3);
186        assert_eq!(slice1.len(), 3);
187        let expected = PrimitiveArray::from_iter(vec![1i32, 1, 1]).into_array();
188        assert_arrays_eq!(slice1, expected);
189
190        // Slice from start of second run to end of second run (indices 3..6)
191        let slice2 = arr.slice(3..6);
192        assert_eq!(slice2.len(), 3);
193        let expected = PrimitiveArray::from_iter(vec![4i32, 4, 4]).into_array();
194        assert_arrays_eq!(slice2, expected);
195
196        // Slice from start of third run to end of third run (indices 6..8)
197        let slice3 = arr.slice(6..8);
198        assert_eq!(slice3.len(), 2);
199        let expected = PrimitiveArray::from_iter(vec![2i32, 2]).into_array();
200        assert_arrays_eq!(slice3, expected);
201
202        // Slice from start of last run to end of last run (indices 8..12)
203        let slice4 = arr.slice(8..12);
204        assert_eq!(slice4.len(), 4);
205        let expected = PrimitiveArray::from_iter(vec![5i32, 5, 5, 5]).into_array();
206        assert_arrays_eq!(slice4, expected);
207
208        // Slice spanning exactly two runs (indices 3..8)
209        let slice5 = arr.slice(3..8);
210        assert_eq!(slice5.len(), 5);
211        let expected = PrimitiveArray::from_iter(vec![4i32, 4, 4, 2, 2]).into_array();
212        assert_arrays_eq!(slice5, expected);
213
214        // Slice from middle of first run to end of second run (indices 1..6)
215        let slice6 = arr.slice(1..6);
216        assert_eq!(slice6.len(), 5);
217        let expected = PrimitiveArray::from_iter(vec![1i32, 1, 4, 4, 4]).into_array();
218        assert_arrays_eq!(slice6, expected);
219
220        // Slice from start of second run to middle of third run (indices 3..7)
221        let slice7 = arr.slice(3..7);
222        assert_eq!(slice7.len(), 4);
223        let expected = PrimitiveArray::from_iter(vec![4i32, 4, 4, 2]).into_array();
224        assert_arrays_eq!(slice7, expected);
225    }
226}
vortex_runend/ops.rs

vortex_runend/
ops.rs
