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

vortex_fastlanes/rle/array/
mod.rs
