vortex_fastlanes/rle/array/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Display;
5use std::fmt::Formatter;
6
7use vortex_array::ArrayRef;
8use vortex_array::ExecutionCtx;
9use vortex_array::TypedArrayRef;
10use vortex_array::array_slots;
11use vortex_error::VortexResult;
12use vortex_error::vortex_ensure;
13
14pub mod rle_compress;
15pub mod rle_decompress;
16
17#[array_slots(crate::RLE)]
18pub struct RLESlots {
19    /// Run values in the dictionary.
20    pub values: ArrayRef,
21    /// Chunk-local indices from all chunks. The start of each chunk is looked up in `values_idx_offsets`.
22    pub indices: ArrayRef,
23    /// Index start positions of each value chunk.
24    ///
25    /// # Example
26    /// ```text
27    /// // Chunk 0: [10, 20] (starts at index 0)
28    /// // Chunk 1: [30, 40] (starts at index 2)
29    /// let values = [10, 20, 30, 40];           // Global values array
30    /// let values_idx_offsets = [0, 2];         // Chunk 0 starts at index 0, Chunk 1 starts at index 2
31    /// ```
32    pub values_idx_offsets: ArrayRef,
33}
34
35#[derive(Clone, Debug)]
36pub struct RLEData {
37    // Offset relative to the start of the chunk.
38    pub(super) offset: usize,
39}
40
41impl Display for RLEData {
42    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
43        write!(f, "offset: {}", self.offset)
44    }
45}
46
47impl RLEData {
48    /// Create a new chunk-based RLE array from its components.
49    ///
50    /// # Arguments
51    ///
52    /// * `values` - Unique values from all chunks
53    /// * `indices` - Chunk-local indices from all chunks
54    /// * `values_idx_offsets` - Start indices for each value chunk.
55    /// * `offset` - Offset into the first chunk
56    /// * `length` - Array length
57    pub fn try_new(offset: usize) -> VortexResult<Self> {
58        vortex_ensure!(
59            offset < 1024,
60            "Offset must be smaller than 1024, got {}",
61            offset
62        );
63        Ok(Self { offset })
64    }
65
66    /// Create a new RLEArray without validation.
67    ///
68    /// # Safety
69    /// The caller must ensure that:
70    /// - `offset + length` does not exceed the length of the indices array
71    /// - The `indices` array contains valid indices into chunks of the `values` array
72    /// - The `values_idx_offsets` array contains valid chunk start offsets
73    pub unsafe fn new_unchecked(offset: usize) -> Self {
74        Self { offset }
75    }
76
77    #[inline]
78    pub fn offset(&self) -> usize {
79        self.offset
80    }
81}
82
83pub trait RLEArrayExt: RLEArraySlotsExt {
84    /// Values index offset relative to the first chunk.
85    ///
86    /// Offsets in `values_idx_offsets` are absolute and need to be shifted
87    /// by the offset of the first chunk, respective the current slice, in
88    /// order to make them relative.
89    #[expect(
90        clippy::expect_used,
91        reason = "expect is safe here as scalar_at returns a valid primitive"
92    )]
93    fn values_idx_offset(&self, chunk_idx: usize, ctx: &mut ExecutionCtx) -> usize {
94        self.values_idx_offsets()
95            .execute_scalar(chunk_idx, ctx)
96            .expect("index must be in bounds")
97            .as_primitive()
98            .as_::<usize>()
99            .expect("index must be of type usize")
100            - self
101                .values_idx_offsets()
102                .execute_scalar(0, ctx)
103                .expect("index must be in bounds")
104                .as_primitive()
105                .as_::<usize>()
106                .expect("index must be of type usize")
107    }
108
109    /// Index offset into the array
110    #[inline]
111    fn offset(&self) -> usize {
112        self.offset
113    }
114}
115
116impl<T: TypedArrayRef<crate::RLE>> RLEArrayExt for T {}
117
118#[cfg(test)]
119mod tests {
120    use vortex_array::ArrayContext;
121    use vortex_array::Canonical;
122    use vortex_array::IntoArray;
123    use vortex_array::VortexSessionExecute;
124    use vortex_array::arrays::PrimitiveArray;
125    use vortex_array::arrays::primitive::PrimitiveArrayExt;
126    use vortex_array::assert_arrays_eq;
127    use vortex_array::dtype::DType;
128    use vortex_array::dtype::Nullability;
129    use vortex_array::dtype::PType;
130    use vortex_array::serde::SerializeOptions;
131    use vortex_array::serde::SerializedArray;
132    use vortex_array::validity::Validity;
133    use vortex_buffer::Buffer;
134    use vortex_buffer::ByteBufferMut;
135    use vortex_error::VortexExpect;
136    use vortex_error::VortexResult;
137    use vortex_session::registry::ReadContext;
138
139    use crate::FL_CHUNK_SIZE;
140    use crate::RLE;
141    use crate::RLEData;
142    use crate::rle::array::RLEArrayExt;
143    use crate::rle::array::RLEArraySlotsExt;
144    use crate::test::SESSION;
145
146    #[test]
147    fn test_try_new() {
148        let values = PrimitiveArray::from_iter([10u32, 20, 30]).into_array();
149
150        // Pad indices to 1024 chunk.
151        let indices =
152            PrimitiveArray::from_iter([0u16, 0, 1, 1, 2].iter().cycle().take(1024).copied())
153                .into_array();
154        let values_idx_offsets = PrimitiveArray::from_iter([0u64]).into_array();
155        let rle_array = RLE::try_new(values, indices, values_idx_offsets, 0, 5)
156            .vortex_expect("RLEData is always valid");
157
158        assert_eq!(rle_array.len(), 5);
159        assert_eq!(rle_array.values().len(), 3);
160        assert_eq!(rle_array.values().dtype().as_ptype(), PType::U32);
161    }
162
163    #[test]
164    fn test_try_new_with_validity() {
165        let values = PrimitiveArray::from_iter([10u32, 20]).into_array();
166        let values_idx_offsets = PrimitiveArray::from_iter([0u64]).into_array();
167
168        let indices_pattern = [0u16, 1, 0];
169        let validity_pattern = [true, false, true];
170
171        // Pad indices to 1024 chunk.
172        let indices_with_validity = PrimitiveArray::new(
173            indices_pattern
174                .iter()
175                .cycle()
176                .take(1024)
177                .copied()
178                .collect::<Buffer<u16>>(),
179            Validity::from_iter(validity_pattern.iter().cycle().take(1024).copied()),
180        )
181        .into_array();
182
183        let rle_array = RLE::try_new(values, indices_with_validity, values_idx_offsets, 0, 3)
184            .vortex_expect("RLEData is always valid");
185
186        assert_eq!(rle_array.len(), 3);
187        assert_eq!(rle_array.values().len(), 2);
188        let mut ctx = SESSION.create_execution_ctx();
189        assert!(rle_array.is_valid(0, &mut ctx).unwrap());
190        assert!(!rle_array.is_valid(1, &mut ctx).unwrap());
191        assert!(rle_array.is_valid(2, &mut ctx).unwrap());
192    }
193
194    #[test]
195    fn test_all_valid() {
196        let values = PrimitiveArray::from_iter([10u32, 20, 30]).into_array();
197        let values_idx_offsets = PrimitiveArray::from_iter([0u64]).into_array();
198
199        let indices_pattern = [0u16, 1, 2, 0, 1];
200        let validity_pattern = [true, true, true, false, false];
201
202        // Pad indices to 1024 chunk.
203        let indices_with_validity = PrimitiveArray::new(
204            indices_pattern
205                .iter()
206                .cycle()
207                .take(1024)
208                .copied()
209                .collect::<Buffer<u16>>(),
210            Validity::from_iter(validity_pattern.iter().cycle().take(1024).copied()),
211        )
212        .into_array();
213
214        let rle_array = RLE::try_new(values, indices_with_validity, values_idx_offsets, 0, 5)
215            .vortex_expect("RLEData is always valid");
216
217        let mut ctx = SESSION.create_execution_ctx();
218        let valid_slice = rle_array
219            .slice(0..3)
220            .unwrap()
221            .execute::<PrimitiveArray>(&mut ctx)
222            .unwrap();
223        // TODO(joe): replace with compute null count
224        assert!(valid_slice.all_valid(&mut ctx).unwrap());
225
226        let mixed_slice = rle_array.slice(1..5).unwrap();
227        assert!(!mixed_slice.all_valid(&mut ctx).unwrap());
228    }
229
230    #[test]
231    fn test_all_invalid() {
232        let values = PrimitiveArray::from_iter([10u32, 20, 30]).into_array();
233        let values_idx_offsets = PrimitiveArray::from_iter([0u64]).into_array();
234
235        // Pad indices to 1024 chunk.
236        let indices_pattern = [0u16, 1, 2, 0, 1];
237        let validity_pattern = [true, true, false, false, false];
238
239        let indices_with_validity = PrimitiveArray::new(
240            indices_pattern
241                .iter()
242                .cycle()
243                .take(1024)
244                .copied()
245                .collect::<Buffer<u16>>(),
246            Validity::from_iter(validity_pattern.iter().cycle().take(1024).copied()),
247        )
248        .into_array();
249
250        let rle_array = RLE::try_new(values, indices_with_validity, values_idx_offsets, 0, 5)
251            .vortex_expect("RLEData is always valid");
252
253        // TODO(joe): replace with compute null count
254        let invalid_slice = rle_array
255            .slice(2..5)
256            .unwrap()
257            .execute::<Canonical>(&mut SESSION.create_execution_ctx())
258            .unwrap()
259            .into_primitive();
260        let mut ctx = SESSION.create_execution_ctx();
261        assert!(invalid_slice.all_invalid(&mut ctx).unwrap());
262
263        let mixed_slice = rle_array.slice(1..4).unwrap();
264        assert!(!mixed_slice.all_invalid(&mut ctx).unwrap());
265    }
266
267    #[test]
268    fn test_validity_mask() {
269        let values = PrimitiveArray::from_iter([10u32, 20, 30]).into_array();
270        let values_idx_offsets = PrimitiveArray::from_iter([0u64]).into_array();
271
272        // Pad indices to 1024 chunk.
273        let indices_pattern = [0u16, 1, 2, 0];
274        let validity_pattern = [true, false, true, false];
275
276        let indices_with_validity = PrimitiveArray::new(
277            indices_pattern
278                .iter()
279                .cycle()
280                .take(1024)
281                .copied()
282                .collect::<Buffer<u16>>(),
283            Validity::from_iter(validity_pattern.iter().cycle().take(1024).copied()),
284        )
285        .into_array();
286
287        let rle_array = RLE::try_new(values, indices_with_validity, values_idx_offsets, 0, 4)
288            .vortex_expect("RLEData is always valid");
289
290        let sliced_array = rle_array.slice(1..4).unwrap();
291        let validity_mask = sliced_array
292            .validity()
293            .unwrap()
294            .execute_mask(sliced_array.len(), &mut SESSION.create_execution_ctx())
295            .unwrap();
296
297        let mut ctx = SESSION.create_execution_ctx();
298        let expected_mask = Validity::from_iter([false, true, false])
299            .execute_mask(3, &mut ctx)
300            .unwrap();
301        assert_eq!(validity_mask.len(), expected_mask.len());
302        assert_eq!(validity_mask, expected_mask);
303        assert_eq!(validity_mask.len(), expected_mask.len());
304        assert_eq!(validity_mask, expected_mask);
305    }
306
307    #[test]
308    fn test_try_new_empty() {
309        let values = PrimitiveArray::from_iter(Vec::<u32>::new()).into_array();
310        let indices = PrimitiveArray::from_iter(Vec::<u16>::new()).into_array();
311        let values_idx_offsets = PrimitiveArray::from_iter(Vec::<u64>::new()).into_array();
312        let rle_array = RLE::try_new(
313            values,
314            indices.clone(),
315            values_idx_offsets,
316            0,
317            indices.len(),
318        )
319        .vortex_expect("RLEData is always valid");
320
321        assert_eq!(rle_array.len(), 0);
322        assert_eq!(rle_array.values().len(), 0);
323    }
324
325    #[test]
326    fn test_multi_chunk_two_chunks() {
327        let mut ctx = SESSION.create_execution_ctx();
328        let values = PrimitiveArray::from_iter([10u32, 20, 30, 40]).into_array();
329        let indices = PrimitiveArray::from_iter([0u16, 1].repeat(1024)).into_array();
330        let values_idx_offsets = PrimitiveArray::from_iter([0u64, 2]).into_array();
331        let rle_array = RLE::try_new(values, indices, values_idx_offsets, 0, 2048)
332            .vortex_expect("RLEData is always valid");
333
334        assert_eq!(rle_array.len(), 2048);
335        assert_eq!(rle_array.values().len(), 4);
336
337        assert_eq!(rle_array.values_idx_offset(0, &mut ctx), 0);
338        assert_eq!(rle_array.values_idx_offset(1, &mut ctx), 2);
339    }
340
341    #[test]
342    fn test_rle_serialization() -> VortexResult<()> {
343        let mut exec_ctx = SESSION.create_execution_ctx();
344        let primitive = PrimitiveArray::from_iter((0..2048).map(|i| (i / 100) as u32));
345        let rle_array = RLEData::encode(primitive.as_view(), &mut exec_ctx)?;
346        assert_eq!(rle_array.len(), 2048);
347
348        let original_data = rle_array
349            .as_array()
350            .clone()
351            .execute::<PrimitiveArray>(&mut exec_ctx)?;
352
353        let ctx = ArrayContext::empty();
354        let serialized =
355            rle_array
356                .into_array()
357                .serialize(&ctx, &SESSION, &SerializeOptions::default())?;
358
359        let mut concat = ByteBufferMut::empty();
360        for buf in serialized {
361            concat.extend_from_slice(buf.as_ref());
362        }
363        let concat = concat.freeze();
364
365        let parts = SerializedArray::try_from(concat)?;
366        let decoded = parts.decode(
367            &DType::Primitive(PType::U32, Nullability::NonNullable),
368            2048,
369            &ReadContext::new(ctx.to_ids()),
370            &SESSION,
371        )?;
372
373        let decoded_data = decoded.execute::<PrimitiveArray>(&mut exec_ctx)?;
374
375        assert_arrays_eq!(
376            original_data,
377            decoded_data,
378            &mut SESSION.create_execution_ctx()
379        );
380        Ok(())
381    }
382
383    #[test]
384    fn test_rle_serialization_slice() -> VortexResult<()> {
385        let mut exec_ctx = SESSION.create_execution_ctx();
386        let primitive = PrimitiveArray::from_iter((0..2048).map(|i| (i / 100) as u32));
387        let rle_array = RLEData::encode(primitive.as_view(), &mut exec_ctx)?;
388
389        let sliced = RLE::try_new(
390            rle_array.values().clone(),
391            rle_array.indices().clone(),
392            rle_array.values_idx_offsets().clone(),
393            100,
394            100,
395        )
396        .vortex_expect("RLEData is always valid");
397        assert_eq!(sliced.len(), 100);
398
399        let ctx = ArrayContext::empty();
400        let serialized =
401            sliced
402                .clone()
403                .into_array()
404                .serialize(&ctx, &SESSION, &SerializeOptions::default())?;
405
406        let mut concat = ByteBufferMut::empty();
407        for buf in serialized {
408            concat.extend_from_slice(buf.as_ref());
409        }
410        let concat = concat.freeze();
411
412        let parts = SerializedArray::try_from(concat)?;
413        let decoded = parts.decode(
414            sliced.dtype(),
415            sliced.len(),
416            &ReadContext::new(ctx.to_ids()),
417            &SESSION,
418        )?;
419
420        let original_data = sliced
421            .as_array()
422            .clone()
423            .execute::<PrimitiveArray>(&mut exec_ctx)?;
424        let decoded_data = decoded.execute::<PrimitiveArray>(&mut exec_ctx)?;
425
426        assert_arrays_eq!(
427            original_data,
428            decoded_data,
429            &mut SESSION.create_execution_ctx()
430        );
431        Ok(())
432    }
433
434    /// Regression test: re-encoding RLE indices with RLE must not corrupt
435    /// chunk-local index values via cross-chunk fill-forward.
436    ///
437    /// The scenario: an array spanning 2 chunks where chunk 0 has 2 distinct
438    /// non-null values (producing chunk-local indices 0 and 1) and chunk 1 is
439    /// entirely null. When fill_forward_nulls propagated the last valid index
440    /// (1) from chunk 0 into chunk 1 during re-encoding, decoding panicked
441    /// because chunk 1 only had 1 unique value and index 1 was out of bounds.
442    #[test]
443    fn test_recompress_indices_no_cross_chunk_leak() -> VortexResult<()> {
444        let mut ctx = SESSION.create_execution_ctx();
445        let len = FL_CHUNK_SIZE + 100;
446        let mut values: Vec<Option<i16>> = vec![None; len];
447        // Two distinct values in chunk 0 → indices 0 and 1.
448        values[0] = Some(10);
449        values[500] = Some(20);
450        // Chunk 1 (positions 1024..) is all-null.
451
452        let original = PrimitiveArray::from_option_iter(values);
453        let rle = RLEData::encode(original.as_view(), &mut ctx)?;
454
455        // Simulate cascading compression: narrow u16->u8 then re-encode with RLE,
456        // matching the path taken by the BtrBlocks compressor.
457        let indices_prim = rle
458            .indices()
459            .clone()
460            .execute::<PrimitiveArray>(&mut ctx)?
461            .narrow(&mut ctx)?;
462        let re_encoded = RLEData::encode(indices_prim.as_view(), &mut ctx)?;
463
464        // Reconstruct the outer RLE with re-encoded indices.
465        // SAFETY: we only replace the indices child; all other invariants hold.
466        let reconstructed = unsafe {
467            RLE::new_unchecked(
468                rle.values().clone(),
469                re_encoded.into_array(),
470                rle.values_idx_offsets().clone(),
471                rle.offset(),
472                rle.len(),
473            )
474        };
475
476        // Decompress — panicked before the fill_forward_nulls chunk-boundary fix.
477        let decoded = reconstructed
478            .as_array()
479            .clone()
480            .execute::<PrimitiveArray>(&mut ctx)?;
481        assert_arrays_eq!(decoded, original, &mut ctx);
482        Ok(())
483    }
484}
vortex_fastlanes/rle/array/mod.rs

vortex_fastlanes/rle/array/
mod.rs
