Skip to main content

vortex_alp/alp/
compress.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use itertools::Itertools;
5use vortex_array::ArrayRef;
6use vortex_array::ArrayView;
7use vortex_array::ExecutionCtx;
8use vortex_array::IntoArray;
9use vortex_array::arrays::Primitive;
10use vortex_array::arrays::PrimitiveArray;
11use vortex_array::dtype::PType;
12use vortex_array::patches::Patches;
13use vortex_array::validity::Validity;
14use vortex_buffer::Buffer;
15use vortex_buffer::BufferMut;
16use vortex_error::VortexResult;
17use vortex_error::vortex_bail;
18use vortex_mask::Mask;
19
20use crate::ALP;
21use crate::Exponents;
22use crate::alp::ALPArray;
23use crate::alp::ALPFloat;
24
25#[macro_export]
26macro_rules! match_each_alp_float_ptype {
27    ($self:expr, | $enc:ident | $body:block) => {{
28        use vortex_array::dtype::PType;
29        use vortex_error::vortex_panic;
30        let ptype = $self;
31        match ptype {
32            PType::F32 => {
33                type $enc = f32;
34                $body
35            }
36            PType::F64 => {
37                type $enc = f64;
38                $body
39            }
40            _ => vortex_panic!("ALP can only encode f32 and f64, got {}", ptype),
41        }
42    }};
43}
44
45pub fn alp_encode(
46    parray: ArrayView<'_, Primitive>,
47    exponents: Option<Exponents>,
48    ctx: &mut ExecutionCtx,
49) -> VortexResult<ALPArray> {
50    let (exponents, encoded, patches) = match parray.ptype() {
51        PType::F32 => alp_encode_components_typed::<f32>(parray, exponents, ctx)?,
52        PType::F64 => alp_encode_components_typed::<f64>(parray, exponents, ctx)?,
53        _ => vortex_bail!("ALP can only encode f32 and f64"),
54    };
55
56    // SAFETY: alp_encode_components_typed must return well-formed components
57    unsafe { Ok(ALP::new_unchecked(encoded, exponents, patches)) }
58}
59
60#[expect(
61    clippy::cast_possible_truncation,
62    reason = "u64 index cast to usize is safe for reasonable array sizes"
63)]
64fn alp_encode_components_typed<T>(
65    values: ArrayView<'_, Primitive>,
66    exponents: Option<Exponents>,
67    ctx: &mut ExecutionCtx,
68) -> VortexResult<(Exponents, ArrayRef, Option<Patches>)>
69where
70    T: ALPFloat,
71{
72    let values_slice = values.as_slice::<T>();
73
74    let (exponents, encoded, exceptional_positions, exceptional_values, mut chunk_offsets) =
75        T::encode(values_slice, exponents);
76
77    let encoded_array = PrimitiveArray::new(encoded, values.validity()?).into_array();
78
79    let validity = values
80        .array()
81        .validity()?
82        .execute_mask(values.array().len(), ctx)?;
83    // exceptional_positions may contain exceptions at invalid positions (which contain garbage
84    // data). We remove null exceptions in order to keep the Patches small.
85    let (valid_exceptional_positions, valid_exceptional_values): (Buffer<u64>, Buffer<T>) =
86        match validity {
87            Mask::AllTrue(_) => (exceptional_positions, exceptional_values),
88            Mask::AllFalse(_) => {
89                // no valid positions, ergo nothing worth patching
90                (Buffer::empty(), Buffer::empty())
91            }
92            Mask::Values(is_valid) => {
93                let (pos, vals): (BufferMut<u64>, BufferMut<T>) = exceptional_positions
94                    .into_iter()
95                    .zip_eq(exceptional_values)
96                    .filter(|(index, _)| {
97                        let is_valid = is_valid.value(*index as usize);
98                        if !is_valid {
99                            let patch_chunk = *index as usize / 1024;
100                            for chunk_idx in (patch_chunk + 1)..chunk_offsets.len() {
101                                chunk_offsets[chunk_idx] -= 1;
102                            }
103                        }
104                        is_valid
105                    })
106                    .unzip();
107                (pos.freeze(), vals.freeze())
108            }
109        };
110    let patches = if valid_exceptional_positions.is_empty() {
111        None
112    } else {
113        let patches_validity = if values.dtype().is_nullable() {
114            Validity::AllValid
115        } else {
116            Validity::NonNullable
117        };
118        let valid_exceptional_values =
119            PrimitiveArray::new(valid_exceptional_values, patches_validity).into_array();
120
121        Some(Patches::new(
122            values_slice.len(),
123            0,
124            valid_exceptional_positions.into_array(),
125            valid_exceptional_values,
126            Some(chunk_offsets.into_array()),
127        )?)
128    };
129    Ok((exponents, encoded_array, patches))
130}
131
132#[cfg(test)]
133mod tests {
134    use core::f32;
135    use core::f64;
136    use std::sync::LazyLock;
137
138    use f64::consts::E;
139    use f64::consts::PI;
140    use vortex_array::VortexSessionExecute;
141    use vortex_array::assert_arrays_eq;
142    use vortex_array::dtype::NativePType;
143    use vortex_array::validity::Validity;
144    use vortex_buffer::Buffer;
145    use vortex_buffer::buffer;
146    use vortex_session::VortexSession;
147
148    use super::*;
149    use crate::alp::array::ALPArrayExt;
150    use crate::alp::array::ALPArraySlotsExt;
151    use crate::decompress_into_array;
152
153    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
154        let session = vortex_array::array_session();
155        crate::initialize(&session);
156        session
157    });
158
159    #[test]
160    fn test_compress() {
161        let array = PrimitiveArray::new(buffer![1.234f32; 1025], Validity::NonNullable);
162        let encoded =
163            alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
164        assert!(encoded.patches().is_none());
165        let expected_encoded = PrimitiveArray::from_iter(vec![1234i32; 1025]);
166        assert_arrays_eq!(
167            encoded.encoded(),
168            expected_encoded,
169            &mut SESSION.create_execution_ctx()
170        );
171        assert_eq!(encoded.exponents(), Exponents { e: 9, f: 6 });
172
173        let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap();
174        assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx());
175    }
176
177    #[test]
178    fn test_nullable_compress() {
179        let array = PrimitiveArray::from_option_iter([None, Some(1.234f32), None]);
180        let encoded =
181            alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
182        assert!(encoded.patches().is_none());
183        let expected_encoded = PrimitiveArray::from_option_iter([None, Some(1234i32), None]);
184        assert_arrays_eq!(
185            encoded.encoded(),
186            expected_encoded,
187            &mut SESSION.create_execution_ctx()
188        );
189        assert_eq!(encoded.exponents(), Exponents { e: 9, f: 6 });
190
191        let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap();
192        let expected = PrimitiveArray::from_option_iter(vec![None, Some(1.234f32), None]);
193        assert_arrays_eq!(decoded, expected, &mut SESSION.create_execution_ctx());
194    }
195
196    #[test]
197    #[expect(clippy::approx_constant)] // Clippy objects to 2.718, an approximation of e, the base of the natural logarithm.
198    fn test_patched_compress() {
199        let values = buffer![1.234f64, 2.718, PI, 4.0];
200        let array = PrimitiveArray::new(values.clone(), Validity::NonNullable);
201        let encoded =
202            alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
203        assert!(encoded.patches().is_some());
204        let expected_encoded = PrimitiveArray::from_iter(vec![1234i64, 2718, 1234, 4000]);
205        assert_arrays_eq!(
206            encoded.encoded(),
207            expected_encoded,
208            &mut SESSION.create_execution_ctx()
209        );
210        assert_eq!(encoded.exponents(), Exponents { e: 16, f: 13 });
211
212        let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap();
213        let expected_decoded = PrimitiveArray::new(values, Validity::NonNullable);
214        assert_arrays_eq!(
215            decoded,
216            expected_decoded,
217            &mut SESSION.create_execution_ctx()
218        );
219    }
220
221    #[test]
222    #[expect(clippy::approx_constant)] // Clippy objects to 2.718, an approximation of e, the base of the natural logarithm.
223    fn test_compress_ignores_invalid_exceptional_values() {
224        let values = buffer![1.234f64, 2.718, PI, 4.0];
225        let array = PrimitiveArray::new(values, Validity::from_iter([true, true, false, true]));
226        let encoded =
227            alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
228        assert!(encoded.patches().is_none());
229        let expected_encoded =
230            PrimitiveArray::from_option_iter(buffer![Some(1234i64), Some(2718), None, Some(4000)]);
231        assert_arrays_eq!(
232            encoded.encoded(),
233            expected_encoded,
234            &mut SESSION.create_execution_ctx()
235        );
236        assert_eq!(encoded.exponents(), Exponents { e: 16, f: 13 });
237
238        let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap();
239        assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx());
240    }
241
242    #[test]
243    #[expect(clippy::approx_constant)] // ALP doesn't like E
244    fn test_nullable_patched_scalar_at() {
245        let array = PrimitiveArray::from_option_iter([
246            Some(1.234f64),
247            Some(2.718),
248            Some(PI),
249            Some(4.0),
250            None,
251        ]);
252        let encoded =
253            alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
254        assert!(encoded.patches().is_some());
255
256        assert_eq!(encoded.exponents(), Exponents { e: 16, f: 13 });
257
258        assert_arrays_eq!(encoded, array, &mut SESSION.create_execution_ctx());
259
260        let _decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap();
261    }
262
263    #[test]
264    fn roundtrips_close_fractional() {
265        let original = PrimitiveArray::from_iter([195.26274f32, 195.27837, -48.815685]);
266        let alp_arr = alp_encode(
267            original.as_view(),
268            None,
269            &mut SESSION.create_execution_ctx(),
270        )
271        .unwrap();
272        assert_arrays_eq!(alp_arr, original, &mut SESSION.create_execution_ctx());
273    }
274
275    #[test]
276    fn roundtrips_all_null() {
277        let mut ctx = SESSION.create_execution_ctx();
278        let original =
279            PrimitiveArray::new(buffer![195.26274f64, PI, -48.815685], Validity::AllInvalid);
280        let alp_arr = alp_encode(original.as_view(), None, &mut ctx).unwrap();
281        let decompressed = alp_arr
282            .into_array()
283            .execute::<PrimitiveArray>(&mut ctx)
284            .unwrap();
285
286        assert_eq!(
287            // The second and third values become exceptions and are replaced
288            [195.26274, 195.26274, 195.26274],
289            decompressed.as_slice::<f64>()
290        );
291
292        assert_arrays_eq!(decompressed, original, &mut ctx);
293    }
294
295    #[test]
296    fn non_finite_numbers() {
297        let mut ctx = SESSION.create_execution_ctx();
298        let original = PrimitiveArray::new(
299            buffer![0.0f32, -0.0, f32::NAN, f32::NEG_INFINITY, f32::INFINITY],
300            Validity::NonNullable,
301        );
302        let encoded = alp_encode(original.as_view(), None, &mut ctx).unwrap();
303        let decoded = encoded
304            .as_array()
305            .clone()
306            .execute::<PrimitiveArray>(&mut ctx)
307            .unwrap();
308        for idx in 0..original.len() {
309            let decoded_val = decoded.as_slice::<f32>()[idx];
310            let original_val = original.as_slice::<f32>()[idx];
311            assert!(
312                decoded_val.is_eq(original_val),
313                "Expected {original_val} but got {decoded_val}"
314            );
315        }
316    }
317
318    #[test]
319    fn test_chunk_offsets() {
320        let mut ctx = SESSION.create_execution_ctx();
321        let mut values = vec![1.0f64; 3072];
322
323        values[1023] = PI;
324        values[1024] = E;
325        values[1025] = PI;
326
327        let array = PrimitiveArray::new(Buffer::from(values), Validity::NonNullable);
328        let encoded = alp_encode(array.as_view(), None, &mut ctx).unwrap();
329        let patches = encoded.patches().unwrap();
330
331        let chunk_offsets = patches
332            .chunk_offsets()
333            .clone()
334            .unwrap()
335            .execute::<PrimitiveArray>(&mut ctx)
336            .unwrap();
337        let expected_offsets = PrimitiveArray::from_iter(vec![0u64, 1, 3]);
338        assert_arrays_eq!(chunk_offsets, expected_offsets, &mut ctx);
339
340        let patch_indices = patches
341            .indices()
342            .clone()
343            .execute::<PrimitiveArray>(&mut ctx)
344            .unwrap();
345        let expected_indices = PrimitiveArray::from_iter(vec![1023u64, 1024, 1025]);
346        assert_arrays_eq!(patch_indices, expected_indices, &mut ctx);
347
348        let patch_values = patches
349            .values()
350            .clone()
351            .execute::<PrimitiveArray>(&mut ctx)
352            .unwrap();
353        let expected_values = PrimitiveArray::from_iter(vec![PI, E, PI]);
354        assert_arrays_eq!(patch_values, expected_values, &mut ctx);
355    }
356
357    #[test]
358    fn test_chunk_offsets_no_patches_in_middle() {
359        let mut ctx = SESSION.create_execution_ctx();
360        let mut values = vec![1.0f64; 3072];
361        values[0] = PI;
362        values[2048] = E;
363
364        let array = PrimitiveArray::new(Buffer::from(values), Validity::NonNullable);
365        let encoded = alp_encode(array.as_view(), None, &mut ctx).unwrap();
366        let patches = encoded.patches().unwrap();
367
368        let chunk_offsets = patches
369            .chunk_offsets()
370            .clone()
371            .unwrap()
372            .execute::<PrimitiveArray>(&mut ctx)
373            .unwrap();
374        let expected_offsets = PrimitiveArray::from_iter(vec![0u64, 1, 1]);
375        assert_arrays_eq!(chunk_offsets, expected_offsets, &mut ctx);
376
377        let patch_indices = patches
378            .indices()
379            .clone()
380            .execute::<PrimitiveArray>(&mut ctx)
381            .unwrap();
382        let expected_indices = PrimitiveArray::from_iter(vec![0u64, 2048]);
383        assert_arrays_eq!(patch_indices, expected_indices, &mut ctx);
384
385        let patch_values = patches
386            .values()
387            .clone()
388            .execute::<PrimitiveArray>(&mut ctx)
389            .unwrap();
390        let expected_values = PrimitiveArray::from_iter(vec![PI, E]);
391        assert_arrays_eq!(patch_values, expected_values, &mut ctx);
392    }
393
394    #[test]
395    fn test_chunk_offsets_trailing_empty_chunks() {
396        let mut ctx = SESSION.create_execution_ctx();
397        let mut values = vec![1.0f64; 3072];
398        values[0] = PI;
399
400        let array = PrimitiveArray::new(Buffer::from(values), Validity::NonNullable);
401        let encoded = alp_encode(array.as_view(), None, &mut ctx).unwrap();
402        let patches = encoded.patches().unwrap();
403
404        let chunk_offsets = patches
405            .chunk_offsets()
406            .clone()
407            .unwrap()
408            .execute::<PrimitiveArray>(&mut ctx)
409            .unwrap();
410        let expected_offsets = PrimitiveArray::from_iter(vec![0u64, 1, 1]);
411        assert_arrays_eq!(chunk_offsets, expected_offsets, &mut ctx);
412
413        let patch_indices = patches
414            .indices()
415            .clone()
416            .execute::<PrimitiveArray>(&mut ctx)
417            .unwrap();
418        let expected_indices = PrimitiveArray::from_iter(vec![0u64]);
419        assert_arrays_eq!(patch_indices, expected_indices, &mut ctx);
420
421        let patch_values = patches
422            .values()
423            .clone()
424            .execute::<PrimitiveArray>(&mut ctx)
425            .unwrap();
426        let expected_values = PrimitiveArray::from_iter(vec![PI]);
427        assert_arrays_eq!(patch_values, expected_values, &mut ctx);
428    }
429
430    #[test]
431    fn test_chunk_offsets_single_chunk() {
432        let mut ctx = SESSION.create_execution_ctx();
433        let mut values = vec![1.0f64; 512];
434        values[0] = PI;
435        values[100] = E;
436
437        let array = PrimitiveArray::new(Buffer::from(values), Validity::NonNullable);
438        let encoded = alp_encode(array.as_view(), None, &mut ctx).unwrap();
439        let patches = encoded.patches().unwrap();
440
441        let chunk_offsets = patches
442            .chunk_offsets()
443            .clone()
444            .unwrap()
445            .execute::<PrimitiveArray>(&mut ctx)
446            .unwrap();
447        let expected_offsets = PrimitiveArray::from_iter(vec![0u64]);
448        assert_arrays_eq!(chunk_offsets, expected_offsets, &mut ctx);
449
450        let patch_indices = patches
451            .indices()
452            .clone()
453            .execute::<PrimitiveArray>(&mut ctx)
454            .unwrap();
455        let expected_indices = PrimitiveArray::from_iter(vec![0u64, 100]);
456        assert_arrays_eq!(patch_indices, expected_indices, &mut ctx);
457
458        let patch_values = patches
459            .values()
460            .clone()
461            .execute::<PrimitiveArray>(&mut ctx)
462            .unwrap();
463        let expected_values = PrimitiveArray::from_iter(vec![PI, E]);
464        assert_arrays_eq!(patch_values, expected_values, &mut ctx);
465    }
466
467    #[test]
468    fn test_slice_half_chunk_f32_roundtrip() {
469        // Create 1024 elements, encode, slice to first 512, then decode
470        let values = vec![1.234f32; 1024];
471        let original = PrimitiveArray::new(Buffer::from(values), Validity::NonNullable);
472        let encoded = alp_encode(
473            original.as_view(),
474            None,
475            &mut SESSION.create_execution_ctx(),
476        )
477        .unwrap();
478
479        let sliced_alp = encoded.slice(512..1024).unwrap();
480
481        let expected_slice = original.slice(512..1024).unwrap();
482        assert_arrays_eq!(
483            sliced_alp,
484            expected_slice,
485            &mut SESSION.create_execution_ctx()
486        );
487    }
488
489    #[test]
490    fn test_slice_half_chunk_f64_roundtrip() {
491        let values = vec![5.678f64; 1024];
492        let original = PrimitiveArray::new(Buffer::from(values), Validity::NonNullable);
493        let encoded = alp_encode(
494            original.as_view(),
495            None,
496            &mut SESSION.create_execution_ctx(),
497        )
498        .unwrap();
499
500        let sliced_alp = encoded.slice(512..1024).unwrap();
501
502        let expected_slice = original.slice(512..1024).unwrap();
503        assert_arrays_eq!(
504            sliced_alp,
505            expected_slice,
506            &mut SESSION.create_execution_ctx()
507        );
508    }
509
510    #[test]
511    fn test_slice_half_chunk_with_patches_roundtrip() {
512        let mut values = vec![1.0f64; 1024];
513        values[100] = PI;
514        values[200] = E;
515        values[600] = 42.42;
516
517        let original = PrimitiveArray::new(Buffer::from(values), Validity::NonNullable);
518        let encoded = alp_encode(
519            original.as_view(),
520            None,
521            &mut SESSION.create_execution_ctx(),
522        )
523        .unwrap();
524
525        let sliced_alp = encoded.slice(512..1024).unwrap();
526
527        let expected_slice = original.slice(512..1024).unwrap();
528        assert_arrays_eq!(
529            sliced_alp,
530            expected_slice,
531            &mut SESSION.create_execution_ctx()
532        );
533        assert!(encoded.patches().is_some());
534    }
535
536    #[test]
537    fn test_slice_across_chunks_with_patches_roundtrip() {
538        let mut values = vec![1.0f64; 2048];
539        values[100] = PI;
540        values[200] = E;
541        values[600] = 42.42;
542        values[800] = 42.42;
543        values[1000] = 42.42;
544        values[1023] = 42.42;
545
546        let original = PrimitiveArray::new(Buffer::from(values), Validity::NonNullable);
547        let encoded = alp_encode(
548            original.as_view(),
549            None,
550            &mut SESSION.create_execution_ctx(),
551        )
552        .unwrap();
553
554        let sliced_alp = encoded.slice(1023..1025).unwrap();
555
556        let expected_slice = original.slice(1023..1025).unwrap();
557        assert_arrays_eq!(
558            sliced_alp,
559            expected_slice,
560            &mut SESSION.create_execution_ctx()
561        );
562        assert!(encoded.patches().is_some());
563    }
564
565    #[test]
566    fn test_slice_half_chunk_nullable_roundtrip() {
567        let mut ctx = SESSION.create_execution_ctx();
568        let values = (0..1024)
569            .map(|i| if i % 3 == 0 { None } else { Some(2.5f32) })
570            .collect::<Vec<_>>();
571
572        let original = PrimitiveArray::from_option_iter(values);
573        let encoded = alp_encode(original.as_view(), None, &mut ctx).unwrap();
574
575        let sliced_alp = encoded.slice(512..1024).unwrap();
576        let decoded = sliced_alp.execute::<PrimitiveArray>(&mut ctx).unwrap();
577
578        let expected_slice = original.slice(512..1024).unwrap();
579        assert_arrays_eq!(decoded, expected_slice, &mut ctx);
580    }
581
582    #[test]
583    fn test_large_f32_array_uniform_values() {
584        let size = 10_000;
585        let array = PrimitiveArray::new(buffer![42.125f32; size], Validity::NonNullable);
586        let encoded =
587            alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
588
589        assert!(encoded.patches().is_none());
590        let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap();
591        assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx());
592    }
593
594    #[test]
595    fn test_large_f64_array_uniform_values() {
596        let size = 50_000;
597        let array = PrimitiveArray::new(buffer![123.456789f64; size], Validity::NonNullable);
598        let encoded =
599            alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
600
601        assert!(encoded.patches().is_none());
602        let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap();
603        assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx());
604    }
605
606    #[test]
607    fn test_large_f32_array_with_patches() {
608        let size = 5_000;
609        let mut values = vec![1.5f32; size];
610        values[100] = f32::consts::PI;
611        values[1500] = f32::consts::E;
612        values[3000] = f32::NEG_INFINITY;
613        values[4500] = f32::INFINITY;
614
615        let array = PrimitiveArray::new(Buffer::from(values), Validity::NonNullable);
616        let encoded =
617            alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
618
619        assert!(encoded.patches().is_some());
620        let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap();
621        assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx());
622    }
623
624    #[test]
625    fn test_large_f64_array_with_patches() {
626        let size = 8_000;
627        let mut values = vec![2.2184f64; size];
628        values[0] = PI;
629        values[1000] = E;
630        values[2000] = f64::NAN;
631        values[3000] = f64::INFINITY;
632        values[4000] = f64::NEG_INFINITY;
633        values[5000] = 0.0;
634        values[6000] = -0.0;
635        values[7000] = 999.999999999;
636
637        let array = PrimitiveArray::new(Buffer::from(values.clone()), Validity::NonNullable);
638        let encoded =
639            alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
640
641        assert!(encoded.patches().is_some());
642        let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap();
643
644        for idx in 0..size {
645            let decoded_val = decoded.as_slice::<f64>()[idx];
646            let original_val = values[idx];
647            assert!(
648                decoded_val.is_eq(original_val),
649                "At index {idx}: Expected {original_val} but got {decoded_val}"
650            );
651        }
652    }
653
654    #[test]
655    fn test_large_nullable_array() {
656        let size = 12_000;
657        let values: Vec<Option<f32>> = (0..size)
658            .map(|i| {
659                if i % 7 == 0 {
660                    None
661                } else {
662                    Some((i as f32) * 0.1)
663                }
664            })
665            .collect();
666
667        let array = PrimitiveArray::from_option_iter(values);
668        let encoded =
669            alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
670        let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap();
671
672        assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx());
673    }
674
675    #[test]
676    fn test_large_mixed_validity_with_patches() {
677        let size = 6_000;
678        let mut values = vec![10.125f64; size];
679
680        values[500] = PI;
681        values[1500] = E;
682        values[2500] = f64::INFINITY;
683        values[3500] = f64::NEG_INFINITY;
684        values[4500] = f64::NAN;
685
686        let validity = Validity::from_iter((0..size).map(|i| !matches!(i, 500 | 2500)));
687
688        let array = PrimitiveArray::new(Buffer::from(values), validity);
689        let encoded =
690            alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
691        let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap();
692
693        assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx());
694    }
695
696    /// Regression test for patch_chunk index-out-of-bounds when slicing a multi-chunk
697    /// ALP array mid-chunk with patches in the trailing chunk.
698    ///
699    /// The bug: chunk_offsets are sliced at chunk granularity (1024-row boundaries)
700    /// but patches indices/values are sliced at element granularity. When a slice ends
701    /// mid-chunk, patches_end_idx could exceed patches_indices.len(), causing OOB panic
702    /// during decompression.
703    #[test]
704    fn test_slice_mid_chunk_with_patches_in_trailing_chunk() {
705        // 3 chunks (3072 elements), patches scattered across all chunks.
706        let mut values = vec![1.0f64; 3072];
707        // Chunk 0 patches (indices 0..1024)
708        values[100] = PI;
709        values[500] = E;
710        // Chunk 1 patches (indices 1024..2048)
711        values[1100] = PI;
712        values[1500] = E;
713        values[1900] = PI;
714        // Chunk 2 patches (indices 2048..3072)
715        values[2100] = PI;
716        values[2500] = E;
717        values[2900] = PI;
718
719        let original = PrimitiveArray::new(Buffer::from(values), Validity::NonNullable);
720        let encoded = alp_encode(
721            original.as_view(),
722            None,
723            &mut SESSION.create_execution_ctx(),
724        )
725        .unwrap();
726        assert!(encoded.patches().is_some());
727
728        // Slice ending mid-chunk-2 (element 2500 is inside chunk 2 = 2048..3072).
729        // This creates a mismatch: chunk_offsets includes the full chunk 2 offset,
730        // but patches_indices only includes patches up to element 2500.
731        let sliced_alp = encoded.slice(0..2500).unwrap();
732        let expected = original.slice(0..2500).unwrap();
733        assert_arrays_eq!(sliced_alp, expected, &mut SESSION.create_execution_ctx());
734
735        // Also test slicing that starts mid-chunk (both start and end mid-chunk).
736        let sliced_alp = encoded.slice(500..2500).unwrap();
737        let expected = original.slice(500..2500).unwrap();
738        assert_arrays_eq!(sliced_alp, expected, &mut SESSION.create_execution_ctx());
739    }
740}