vortex_fastlanes/rle/vtable/
operations.rs

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

vortex_fastlanes/rle/vtable/
operations.rs
