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

vortex_runend/
ops.rs
