vortex_fastlanes/rle/vtable/
operations.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::ops::Range;
5
6use vortex_array::ArrayRef;
7use vortex_array::IntoArray;
8use vortex_array::vtable::OperationsVTable;
9use vortex_error::VortexExpect;
10use vortex_scalar::Scalar;
11
12use super::RLEVTable;
13use crate::FL_CHUNK_SIZE;
14use crate::RLEArray;
15
16impl OperationsVTable<RLEVTable> for RLEVTable {
17    fn slice(array: &RLEArray, range: Range<usize>) -> ArrayRef {
18        let offset_in_chunk = array.offset();
19        let chunk_start_idx = (offset_in_chunk + range.start) / FL_CHUNK_SIZE;
20        let chunk_end_idx = (offset_in_chunk + range.end).div_ceil(FL_CHUNK_SIZE);
21
22        let values_start_idx = array.values_idx_offset(chunk_start_idx);
23        let values_end_idx = if chunk_end_idx < array.values_idx_offsets().len() {
24            array.values_idx_offset(chunk_end_idx)
25        } else {
26            array.values().len()
27        };
28
29        let sliced_values = array.values().slice(values_start_idx..values_end_idx);
30
31        let sliced_values_idx_offsets = array
32            .values_idx_offsets()
33            .slice(chunk_start_idx..chunk_end_idx);
34
35        let sliced_indices = array
36            .indices()
37            .slice(chunk_start_idx * FL_CHUNK_SIZE..chunk_end_idx * FL_CHUNK_SIZE);
38
39        // SAFETY: Slicing preserves all invariants.
40        unsafe {
41            RLEArray::new_unchecked(
42                sliced_values,
43                sliced_indices,
44                sliced_values_idx_offsets,
45                array.dtype().clone(),
46                // Keep the offset relative to the first chunk.
47                (array.offset() + range.start) % FL_CHUNK_SIZE,
48                range.len(),
49            )
50            .into_array()
51        }
52    }
53
54    fn scalar_at(array: &RLEArray, index: usize) -> Scalar {
55        let offset_in_chunk = array.offset();
56        let chunk_relative_idx = array.indices().scalar_at(offset_in_chunk + index);
57
58        let chunk_relative_idx = chunk_relative_idx
59            .as_primitive()
60            .as_::<usize>()
61            .vortex_expect("Index must not be null");
62
63        let chunk_id = (offset_in_chunk + index) / FL_CHUNK_SIZE;
64        let value_idx_offset = array.values_idx_offset(chunk_id);
65
66        let scalar = array
67            .values()
68            .scalar_at(value_idx_offset + chunk_relative_idx);
69
70        Scalar::new(array.dtype().clone(), scalar.into_value())
71    }
72}
73
74#[cfg(test)]
75mod tests {
76    use vortex_array::Array;
77    use vortex_array::IntoArray;
78    use vortex_array::ToCanonical;
79    use vortex_array::arrays::PrimitiveArray;
80    use vortex_array::assert_arrays_eq;
81    use vortex_array::validity::Validity;
82    use vortex_buffer::Buffer;
83    use vortex_buffer::buffer;
84
85    use super::*;
86
87    mod fixture {
88        use super::*;
89
90        pub(super) fn rle_array() -> RLEArray {
91            let values = PrimitiveArray::from_iter([10u32, 20u32, 30u32]).into_array();
92            let indices = PrimitiveArray::from_iter(
93                [0u16, 0u16, 1u16, 1u16, 1u16, 2u16, 0u16]
94                    .iter()
95                    .cycle()
96                    .take(1024)
97                    .copied(),
98            )
99            .into_array();
100            let values_idx_offsets = PrimitiveArray::from_iter([0u64]).into_array();
101
102            RLEArray::try_new(
103                values,
104                indices.clone(),
105                values_idx_offsets,
106                0,
107                indices.len(),
108            )
109            .unwrap()
110        }
111
112        pub(super) fn rle_array_with_nulls() -> RLEArray {
113            let values = PrimitiveArray::from_iter([10u32, 20u32, 30u32]).into_array();
114            let pattern = [0u16, 0u16, 1u16, 1u16, 1u16, 2u16, 0u16];
115            let values_idx_offsets = PrimitiveArray::from_iter([0u64]).into_array();
116
117            // Repeat the validity pattern to match indices length
118            let validity = Validity::from_iter(
119                [true, false, true, true, false, true, true]
120                    .iter()
121                    .cycle()
122                    .take(1024)
123                    .copied(),
124            );
125
126            let indices = PrimitiveArray::new(
127                pattern
128                    .iter()
129                    .cycle()
130                    .take(1024)
131                    .copied()
132                    .collect::<Buffer<u16>>(),
133                validity,
134            )
135            .into_array();
136
137            RLEArray::try_new(
138                values,
139                indices.clone(),
140                values_idx_offsets,
141                0,
142                indices.len(),
143            )
144            .unwrap()
145        }
146    }
147
148    #[test]
149    fn test_scalar_at() {
150        use vortex_array::assert_arrays_eq;
151
152        let array = fixture::rle_array();
153        let expected = PrimitiveArray::from_iter([10u32, 10, 20, 20, 20, 30, 10]);
154        assert_arrays_eq!(array.slice(0..7), expected);
155    }
156
157    #[test]
158    fn test_scalar_at_with_nulls() {
159        use vortex_array::assert_arrays_eq;
160
161        let array = fixture::rle_array_with_nulls();
162        let expected = PrimitiveArray::from_option_iter([
163            Some(10u32),
164            None,
165            Some(20),
166            Some(20),
167            None,
168            Some(30),
169            Some(10),
170        ]);
171        assert_arrays_eq!(array.slice(0..7), expected);
172    }
173
174    #[test]
175    fn test_scalar_at_slice() {
176        use vortex_array::assert_arrays_eq;
177
178        let array = fixture::rle_array();
179        let sliced = array.slice(2..6); // [20, 20, 20, 30]
180
181        assert_eq!(sliced.len(), 4);
182        let expected = PrimitiveArray::from_iter([20u32, 20, 20, 30]);
183        assert_arrays_eq!(sliced, expected);
184    }
185
186    #[test]
187    fn test_scalar_at_slice_with_nulls() {
188        use vortex_array::assert_arrays_eq;
189
190        let array = fixture::rle_array_with_nulls();
191        let sliced = array.slice(2..6); // [20, 20, null, 30]
192
193        assert_eq!(sliced.len(), 4);
194        let expected = PrimitiveArray::from_option_iter([Some(20u32), Some(20), None, Some(30)]);
195        assert_arrays_eq!(sliced, expected);
196    }
197
198    #[test]
199    fn test_scalar_at_multiple_chunks() {
200        // Test accessing elements around chunk boundaries
201        let values: Buffer<u16> = (0..3000).map(|i| (i / 50) as u16).collect();
202        let expected: Vec<u16> = (0..3000).map(|i| (i / 50) as u16).collect();
203        let array = values.into_array();
204
205        let encoded = RLEArray::encode(&array.to_primitive()).unwrap();
206
207        // Access scalars from multiple chunks.
208        for &idx in &[1023, 1024, 1025, 2047, 2048, 2049] {
209            if idx < encoded.len() {
210                let original_value = expected[idx];
211                let encoded_value = encoded.scalar_at(idx).as_primitive().as_::<u16>().unwrap();
212                assert_eq!(original_value, encoded_value, "Mismatch at index {}", idx);
213            }
214        }
215    }
216
217    #[test]
218    #[should_panic]
219    fn test_scalar_at_out_of_bounds() {
220        let array = fixture::rle_array();
221        array.scalar_at(1025);
222    }
223
224    #[test]
225    #[should_panic]
226    fn test_scalar_at_slice_out_of_bounds() {
227        let array = fixture::rle_array().slice(0..1);
228        array.scalar_at(1);
229    }
230
231    #[test]
232    fn test_slice_full_range() {
233        let array = fixture::rle_array_with_nulls();
234        let sliced = array.slice(0..7);
235
236        let expected = PrimitiveArray::from_option_iter([
237            Some(10u32),
238            None,
239            Some(20),
240            Some(20),
241            None,
242            Some(30),
243            Some(10),
244        ]);
245        assert_arrays_eq!(sliced.to_array(), expected.to_array());
246    }
247
248    #[test]
249    fn test_slice_partial_range() {
250        let array = fixture::rle_array();
251        let sliced = array.slice(4..6); // [20, 30]
252
253        let expected = buffer![20u32, 30].into_array();
254        assert_arrays_eq!(sliced.to_array(), expected);
255    }
256
257    #[test]
258    fn test_slice_single_element() {
259        let array = fixture::rle_array();
260        let sliced = array.slice(5..6); // [30]
261
262        let expected = buffer![30u32].into_array();
263        assert_arrays_eq!(sliced.to_array(), expected);
264    }
265
266    #[test]
267    fn test_slice_empty_range() {
268        let array = fixture::rle_array();
269        let sliced = array.slice(3..3);
270
271        assert_eq!(sliced.len(), 0);
272    }
273
274    #[test]
275    fn test_slice_with_nulls() {
276        let array = fixture::rle_array_with_nulls();
277        let sliced = array.slice(1..4); // [null, 20, 20]
278
279        let expected = PrimitiveArray::from_option_iter([Option::<u32>::None, Some(20), Some(20)]);
280        assert_arrays_eq!(sliced.to_array(), expected.to_array());
281    }
282
283    #[test]
284    fn test_slice_decode_with_nulls() {
285        let array = fixture::rle_array_with_nulls();
286        let sliced = array.slice(1..4).to_array().to_primitive(); // [null, 20, 20]
287
288        let expected = PrimitiveArray::from_option_iter([Option::<u32>::None, Some(20), Some(20)]);
289        assert_arrays_eq!(sliced.to_array(), expected.to_array());
290    }
291
292    #[test]
293    fn test_slice_preserves_dtype() {
294        let array = fixture::rle_array();
295        let sliced = array.slice(1..4);
296
297        assert_eq!(array.dtype(), sliced.dtype());
298    }
299
300    #[test]
301    fn test_slice_across_chunk_boundaries() {
302        let values: Buffer<u32> = (0..2100).map(|i| (i / 100) as u32).collect();
303        let expected: Vec<u32> = (0..2100).map(|i| (i / 100) as u32).collect();
304        let array = values.into_array();
305
306        let encoded = RLEArray::encode(&array.to_primitive()).unwrap();
307
308        // Slice across first and second chunk.
309        let slice = encoded.slice(500..1500);
310        assert_arrays_eq!(
311            slice,
312            PrimitiveArray::from_iter(expected[500..1500].iter().copied())
313        );
314
315        // Slice across second and third chunk.
316        let slice = encoded.slice(1000..2000);
317        assert_arrays_eq!(
318            slice,
319            PrimitiveArray::from_iter(expected[1000..2000].iter().copied())
320        );
321    }
322}
vortex_fastlanes/rle/vtable/operations.rs

vortex_fastlanes/rle/vtable/
operations.rs
