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

vortex_runend/
ops.rs
