vortex_runend/
ops.rs

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

vortex_runend/
ops.rs
