vortex_runend/
ops.rs

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

vortex_runend/
ops.rs
