vortex_fastlanes/rle/vtable/
operations.rs

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

vortex_fastlanes/rle/vtable/
operations.rs
