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

vortex_runend/
ops.rs
