1use std::fmt::Debug;
5use std::fmt::Display;
6use std::fmt::Formatter;
7use std::hash::Hash;
8use std::hash::Hasher;
9
10use prost::Message;
11use vortex_array::Array;
12use vortex_array::ArrayEq;
13use vortex_array::ArrayHash;
14use vortex_array::ArrayId;
15use vortex_array::ArrayParts;
16use vortex_array::ArrayRef;
17use vortex_array::ArraySlots;
18use vortex_array::ArrayView;
19use vortex_array::EqMode;
20use vortex_array::ExecutionCtx;
21use vortex_array::ExecutionResult;
22use vortex_array::IntoArray;
23use vortex_array::TypedArrayRef;
24use vortex_array::array_slots;
25use vortex_array::arrays::Primitive;
26use vortex_array::buffer::BufferHandle;
27use vortex_array::dtype::DType;
28use vortex_array::dtype::NativePType;
29use vortex_array::dtype::PType;
30use vortex_array::patches::PatchSlotIndices;
31use vortex_array::patches::Patches;
32use vortex_array::patches::PatchesData;
33use vortex_array::patches::PatchesMetadata;
34use vortex_array::require_child;
35use vortex_array::require_patches;
36use vortex_array::serde::ArrayChildren;
37use vortex_array::smallvec::smallvec;
38use vortex_array::vtable::VTable;
39use vortex_array::vtable::ValidityChild;
40use vortex_array::vtable::ValidityVTableFromChild;
41use vortex_error::VortexExpect;
42use vortex_error::VortexResult;
43use vortex_error::vortex_bail;
44use vortex_error::vortex_ensure;
45use vortex_error::vortex_ensure_eq;
46use vortex_error::vortex_panic;
47use vortex_session::VortexSession;
48use vortex_session::registry::CachedId;
49
50use crate::ALPFloat;
51use crate::alp::Exponents;
52use crate::alp::decompress::execute_decompress;
53use crate::alp::rules::RULES;
54
55pub type ALPArray = Array<ALP>;
57
58impl ArrayHash for ALPData {
59 fn array_hash<H: Hasher>(&self, state: &mut H, _accuracy: EqMode) {
60 self.exponents.hash(state);
61 self.patches_data.hash(state);
62 }
63}
64
65impl ArrayEq for ALPData {
66 fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool {
67 self.exponents == other.exponents && self.patches_data == other.patches_data
68 }
69}
70
71impl VTable for ALP {
72 type TypedArrayData = ALPData;
73
74 type OperationsVTable = Self;
75 type ValidityVTable = ValidityVTableFromChild;
76
77 fn id(&self) -> ArrayId {
78 static ID: CachedId = CachedId::new("vortex.alp");
79 *ID
80 }
81
82 fn validate(
83 &self,
84 data: &ALPData,
85 dtype: &DType,
86 len: usize,
87 slots: &[Option<ArrayRef>],
88 ) -> VortexResult<()> {
89 let alp_slots = ALPSlotsView::from_slots(slots);
90 let patches =
91 PatchesData::patches_from_slots(data.patches_data.as_ref(), len, slots, PATCH_SLOTS);
92 validate_parts(dtype, len, data.exponents, alp_slots.encoded, patches)
93 }
94
95 fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
96 0
97 }
98
99 fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
100 vortex_panic!("ALPArray buffer index {idx} out of bounds")
101 }
102
103 fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option<String> {
104 None
105 }
106
107 fn with_buffers(
108 &self,
109 array: ArrayView<'_, Self>,
110 buffers: &[BufferHandle],
111 ) -> VortexResult<ArrayParts<Self>> {
112 vortex_array::vtable::with_empty_buffers(self, array, buffers)
113 }
114
115 fn serialize(
116 array: ArrayView<'_, Self>,
117 _session: &VortexSession,
118 ) -> VortexResult<Option<Vec<u8>>> {
119 let exponents = array.exponents();
120 Ok(Some(
121 ALPMetadata {
122 exp_e: exponents.e as u32,
123 exp_f: exponents.f as u32,
124 patches: array
125 .patches()
126 .map(|p| p.to_metadata(array.len(), array.dtype()))
127 .transpose()?,
128 }
129 .encode_to_vec(),
130 ))
131 }
132
133 fn deserialize(
134 &self,
135 dtype: &DType,
136 len: usize,
137 metadata: &[u8],
138 _buffers: &[BufferHandle],
139 children: &dyn ArrayChildren,
140 _session: &VortexSession,
141 ) -> VortexResult<ArrayParts<Self>> {
142 let metadata = ALPMetadata::decode(metadata)?;
143 let encoded_ptype = match &dtype {
144 DType::Primitive(PType::F32, n) => DType::Primitive(PType::I32, *n),
145 DType::Primitive(PType::F64, n) => DType::Primitive(PType::I64, *n),
146 d => vortex_bail!(MismatchedTypes: "f32 or f64", d),
147 };
148 let encoded = children.get(0, &encoded_ptype, len)?;
149
150 let patches = metadata
151 .patches
152 .map(|p| {
153 let indices = children.get(1, &p.indices_dtype()?, p.len()?)?;
154 let values = children.get(2, dtype, p.len()?)?;
155 let chunk_offsets = p
156 .chunk_offsets_dtype()?
157 .map(|dtype| children.get(3, &dtype, usize::try_from(p.chunk_offsets_len())?))
158 .transpose()?;
159
160 Patches::new(len, p.offset()?, indices, values, chunk_offsets)
161 })
162 .transpose()?;
163
164 let slots = ALPData::make_slots(&encoded, patches.as_ref());
165 let data = ALPData::new(
166 Exponents {
167 e: u8::try_from(metadata.exp_e)?,
168 f: u8::try_from(metadata.exp_f)?,
169 },
170 patches,
171 );
172 Ok(ArrayParts::new(
173 self.clone(),
174 dtype.clone(),
175 len,
176 data,
177 slots,
178 ))
179 }
180
181 fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
182 ALPSlots::NAMES[idx].to_string()
183 }
184
185 fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
186 let array = require_child!(array, array.encoded(), ALPSlots::ENCODED => Primitive);
187 require_patches!(
188 array,
189 ALPSlots::PATCH_INDICES,
190 ALPSlots::PATCH_VALUES,
191 ALPSlots::PATCH_CHUNK_OFFSETS
192 );
193
194 Ok(ExecutionResult::done(
195 execute_decompress(array, ctx)?.into_array(),
196 ))
197 }
198
199 fn reduce_parent(
200 array: ArrayView<'_, Self>,
201 parent: &ArrayRef,
202 child_idx: usize,
203 ) -> VortexResult<Option<ArrayRef>> {
204 RULES.evaluate(array, parent, child_idx)
205 }
206}
207
208#[array_slots(ALP)]
209pub struct ALPSlots {
210 #[slot(0)]
212 pub encoded: ArrayRef,
213 #[slot(1)]
215 pub patch_indices: Option<ArrayRef>,
216 #[slot(2)]
218 pub patch_values: Option<ArrayRef>,
219 #[slot(3)]
221 pub patch_chunk_offsets: Option<ArrayRef>,
222}
223
224const PATCH_SLOTS: PatchSlotIndices = PatchSlotIndices {
225 indices: ALPSlots::PATCH_INDICES,
226 values: ALPSlots::PATCH_VALUES,
227 chunk_offsets: ALPSlots::PATCH_CHUNK_OFFSETS,
228};
229
230#[derive(Clone, Debug)]
231pub struct ALPData {
232 patches_data: Option<PatchesData>,
233 exponents: Exponents,
234}
235
236impl Display for ALPData {
237 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
238 write!(f, "exponents: {}", self.exponents)?;
239 if let Some(pd) = &self.patches_data {
240 write!(f, ", patch_offset: {}", pd.offset())?;
241 }
242 Ok(())
243 }
244}
245
246#[derive(Clone, Debug)]
247pub struct ALP;
248
249#[derive(Clone, prost::Message)]
250pub struct ALPMetadata {
251 #[prost(uint32, tag = "1")]
252 pub(crate) exp_e: u32,
253 #[prost(uint32, tag = "2")]
254 pub(crate) exp_f: u32,
255 #[prost(message, optional, tag = "3")]
256 pub(crate) patches: Option<PatchesMetadata>,
257}
258
259impl ALPData {
260 fn validate_components(
261 encoded: &ArrayRef,
262 exponents: Exponents,
263 patches: Option<&Patches>,
264 ) -> VortexResult<()> {
265 vortex_ensure!(
266 matches!(
267 encoded.dtype(),
268 DType::Primitive(PType::I32 | PType::I64, _)
269 ),
270 "ALP encoded ints have invalid DType {}",
271 encoded.dtype(),
272 );
273
274 let Exponents { e, f } = exponents;
277 match encoded.dtype().as_ptype() {
278 PType::I32 => {
279 vortex_ensure!(exponents.e <= f32::MAX_EXPONENT, "e out of bounds: {e}");
280 vortex_ensure!(exponents.f <= f32::MAX_EXPONENT, "f out of bounds: {f}");
281 if let Some(patches) = patches {
282 Self::validate_patches::<f32>(patches, encoded)?;
283 }
284 }
285 PType::I64 => {
286 vortex_ensure!(e <= f64::MAX_EXPONENT, "e out of bounds: {e}");
287 vortex_ensure!(f <= f64::MAX_EXPONENT, "f out of bounds: {f}");
288
289 if let Some(patches) = patches {
290 Self::validate_patches::<f64>(patches, encoded)?;
291 }
292 }
293 _ => unreachable!(),
294 }
295
296 if let Some(patches) = patches {
298 vortex_ensure_eq!(patches.array_len(), encoded.len());
299
300 }
302
303 Ok(())
304 }
305
306 fn logical_dtype(encoded: &ArrayRef) -> VortexResult<DType> {
307 match encoded.dtype() {
308 DType::Primitive(PType::I32, nullability) => {
309 Ok(DType::Primitive(PType::F32, *nullability))
310 }
311 DType::Primitive(PType::I64, nullability) => {
312 Ok(DType::Primitive(PType::F64, *nullability))
313 }
314 _ => vortex_bail!("ALP encoded ints have invalid DType {}", encoded.dtype(),),
315 }
316 }
317
318 fn validate_patches<T: ALPFloat + NativePType>(
320 patches: &Patches,
321 encoded: &ArrayRef,
322 ) -> VortexResult<()> {
323 vortex_ensure_eq!(patches.array_len(), encoded.len());
324
325 let expected_type = DType::Primitive(T::PTYPE, encoded.dtype().nullability());
326 vortex_ensure_eq!(patches.dtype(), &expected_type);
327
328 Ok(())
329 }
330}
331
332impl ALPData {
333 pub fn new(exponents: Exponents, patches: Option<Patches>) -> Self {
338 Self {
339 patches_data: patches.as_ref().map(PatchesData::from_patches),
340 exponents,
341 }
342 }
343
344 pub(crate) unsafe fn new_unchecked(exponents: Exponents, patches: Option<Patches>) -> Self {
355 Self::new(exponents, patches)
356 }
357}
358
359impl ALP {
361 pub fn new(encoded: ArrayRef, exponents: Exponents, patches: Option<Patches>) -> ALPArray {
362 let dtype = ALPData::logical_dtype(&encoded).vortex_expect("ALP encoded dtype");
363 let len = encoded.len();
364 let slots = ALPData::make_slots(&encoded, patches.as_ref());
365 unsafe {
366 Array::from_parts_unchecked(ArrayParts::new(
367 ALP,
368 dtype,
369 len,
370 ALPData::new(exponents, patches),
371 slots,
372 ))
373 }
374 }
375
376 pub fn try_new(
377 encoded: ArrayRef,
378 exponents: Exponents,
379 patches: Option<Patches>,
380 ) -> VortexResult<ALPArray> {
381 let dtype = ALPData::logical_dtype(&encoded)?;
382 let len = encoded.len();
383 let slots = ALPData::make_slots(&encoded, patches.as_ref());
384 let data = ALPData::new(exponents, patches);
385 Array::try_from_parts(ArrayParts::new(ALP, dtype, len, data, slots))
386 }
387
388 pub unsafe fn new_unchecked(
391 encoded: ArrayRef,
392 exponents: Exponents,
393 patches: Option<Patches>,
394 ) -> ALPArray {
395 let dtype = ALPData::logical_dtype(&encoded).vortex_expect("ALP encoded dtype");
396 let len = encoded.len();
397 let slots = ALPData::make_slots(&encoded, patches.as_ref());
398 let data = unsafe { ALPData::new_unchecked(exponents, patches) };
399 unsafe { Array::from_parts_unchecked(ArrayParts::new(ALP, dtype, len, data, slots)) }
400 }
401}
402
403impl ALPData {
404 fn make_slots(encoded: &ArrayRef, patches: Option<&Patches>) -> ArraySlots {
405 let mut slots: ArraySlots = smallvec![Some(encoded.clone())];
406 PatchesData::push_slots(&mut slots, patches);
407 slots
408 }
409
410 #[inline]
411 pub fn exponents(&self) -> Exponents {
412 self.exponents
413 }
414}
415
416pub trait ALPArrayExt: ALPArraySlotsExt {
417 fn exponents(&self) -> Exponents {
418 self.exponents
419 }
420
421 fn patches(&self) -> Option<Patches> {
422 PatchesData::patches_from_slots(
423 self.patches_data.as_ref(),
424 self.len(),
425 self.slots(),
426 PATCH_SLOTS,
427 )
428 }
429}
430
431fn validate_parts(
432 dtype: &DType,
433 len: usize,
434 exponents: Exponents,
435 encoded: &ArrayRef,
436 patches: Option<Patches>,
437) -> VortexResult<()> {
438 let logical_dtype = ALPData::logical_dtype(encoded)?;
439 ALPData::validate_components(encoded, exponents, patches.as_ref())?;
440 vortex_ensure_eq!(encoded.len(), len, "ALP encoded len must match outer len");
441 vortex_ensure_eq!(
442 &logical_dtype,
443 dtype,
444 "ALP dtype does not match encoded logical dtype"
445 );
446 Ok(())
447}
448
449impl<T: TypedArrayRef<ALP>> ALPArrayExt for T {}
450
451pub trait ALPArrayOwnedExt {
452 fn into_parts(self) -> (ArrayRef, Exponents, Option<Patches>);
453}
454
455impl ALPArrayOwnedExt for Array<ALP> {
456 #[inline]
457 fn into_parts(self) -> (ArrayRef, Exponents, Option<Patches>) {
458 let patches = self.patches();
459 let exponents = self.exponents();
460 let encoded = self.encoded().clone();
461 (encoded, exponents, patches)
462 }
463}
464
465impl ValidityChild<ALP> for ALP {
466 fn validity_child(array: ArrayView<'_, ALP>) -> ArrayRef {
467 array.encoded().clone()
468 }
469}
470
471#[cfg(test)]
472mod tests {
473 use std::f64::consts::PI;
474 use std::sync::LazyLock;
475
476 use rstest::rstest;
477 use vortex_array::VortexSessionExecute;
478 use vortex_array::arrays::PrimitiveArray;
479 use vortex_array::assert_arrays_eq;
480 use vortex_session::VortexSession;
481
482 use super::*;
483 use crate::alp_encode;
484 use crate::decompress_into_array;
485
486 static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
487 let session = vortex_array::array_session();
488 crate::initialize(&session);
489 session
490 });
491
492 #[rstest]
493 #[case(0)]
494 #[case(1)]
495 #[case(100)]
496 #[case(1023)]
497 #[case(1024)]
498 #[case(1025)]
499 #[case(2047)]
500 #[case(2048)]
501 #[case(2049)]
502 fn test_execute_f32(#[case] size: usize) {
503 let mut ctx = SESSION.create_execution_ctx();
504 let values = PrimitiveArray::from_iter((0..size).map(|i| i as f32));
505 let encoded = alp_encode(values.as_view(), None, &mut ctx).unwrap();
506
507 let expected =
509 decompress_into_array(encoded.clone(), &mut SESSION.create_execution_ctx()).unwrap();
510
511 assert_arrays_eq!(expected, values, &mut ctx);
512 assert_arrays_eq!(encoded, values, &mut ctx);
513 }
514
515 #[rstest]
516 #[case(0)]
517 #[case(1)]
518 #[case(100)]
519 #[case(1023)]
520 #[case(1024)]
521 #[case(1025)]
522 #[case(2047)]
523 #[case(2048)]
524 #[case(2049)]
525 fn test_execute_f64(#[case] size: usize) {
526 let values = PrimitiveArray::from_iter((0..size).map(|i| i as f64));
527 let encoded =
528 alp_encode(values.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
529
530 let mut ctx = SESSION.create_execution_ctx();
531 let expected =
533 decompress_into_array(encoded.clone(), &mut SESSION.create_execution_ctx()).unwrap();
534
535 assert_arrays_eq!(expected, values, &mut ctx);
536 assert_arrays_eq!(encoded, values, &mut ctx);
537 }
538
539 #[rstest]
540 #[case(100)]
541 #[case(1023)]
542 #[case(1024)]
543 #[case(1025)]
544 #[case(2047)]
545 #[case(2048)]
546 #[case(2049)]
547 fn test_execute_with_patches(#[case] size: usize) {
548 let values: Vec<f64> = (0..size)
549 .map(|i| match i % 4 {
550 0..=2 => 1.0,
551 _ => PI,
552 })
553 .collect();
554
555 let array = PrimitiveArray::from_iter(values);
556 let encoded =
557 alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
558 assert!(encoded.patches().unwrap().array_len() > 0);
559
560 let mut ctx = SESSION.create_execution_ctx();
561 let expected =
563 decompress_into_array(encoded.clone(), &mut SESSION.create_execution_ctx()).unwrap();
564
565 assert_arrays_eq!(expected, array, &mut ctx);
566 assert_arrays_eq!(encoded, array, &mut ctx);
567 }
568
569 #[rstest]
570 #[case(0)]
571 #[case(1)]
572 #[case(100)]
573 #[case(1023)]
574 #[case(1024)]
575 #[case(1025)]
576 #[case(2047)]
577 #[case(2048)]
578 #[case(2049)]
579 fn test_execute_with_validity(#[case] size: usize) {
580 let values: Vec<Option<f32>> = (0..size)
581 .map(|i| if i % 2 == 1 { None } else { Some(1.0) })
582 .collect();
583
584 let array = PrimitiveArray::from_option_iter(values);
585 let encoded =
586 alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
587
588 let mut ctx = SESSION.create_execution_ctx();
589 let expected =
591 decompress_into_array(encoded.clone(), &mut SESSION.create_execution_ctx()).unwrap();
592
593 assert_arrays_eq!(expected, array, &mut ctx);
594 assert_arrays_eq!(encoded, array, &mut ctx);
595 }
596
597 #[rstest]
598 #[case(100)]
599 #[case(1023)]
600 #[case(1024)]
601 #[case(1025)]
602 #[case(2047)]
603 #[case(2048)]
604 #[case(2049)]
605 fn test_execute_with_patches_and_validity(#[case] size: usize) {
606 let values: Vec<Option<f64>> = (0..size)
607 .map(|idx| match idx % 3 {
608 0 => Some(1.0),
609 1 => None,
610 _ => Some(PI),
611 })
612 .collect();
613
614 let array = PrimitiveArray::from_option_iter(values);
615 let encoded =
616 alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
617 assert!(encoded.patches().unwrap().array_len() > 0);
618
619 let mut ctx = SESSION.create_execution_ctx();
620 let expected =
622 decompress_into_array(encoded.clone(), &mut SESSION.create_execution_ctx()).unwrap();
623
624 assert_arrays_eq!(expected, array, &mut ctx);
625 assert_arrays_eq!(encoded, array, &mut ctx);
626 }
627
628 #[rstest]
629 #[case(500, 100)]
630 #[case(1000, 200)]
631 #[case(2048, 512)]
632 fn test_execute_sliced_vector(#[case] size: usize, #[case] slice_start: usize) {
633 let values: Vec<Option<f64>> = (0..size)
634 .map(|i| {
635 if i % 5 == 0 {
636 None
637 } else if i % 4 == 3 {
638 Some(PI)
639 } else {
640 Some(1.0)
641 }
642 })
643 .collect();
644
645 let mut ctx = SESSION.create_execution_ctx();
646 let array = PrimitiveArray::from_option_iter(values.clone());
647 let encoded = alp_encode(array.as_view(), None, &mut ctx).unwrap();
648
649 let slice_end = size - slice_start;
650 let sliced_encoded = encoded.slice(slice_start..slice_end).unwrap();
651
652 let expected =
653 PrimitiveArray::from_option_iter(values[slice_start..slice_end].iter().copied());
654 assert_arrays_eq!(sliced_encoded, expected, &mut ctx);
655 }
656
657 #[rstest]
658 #[case(500, 100)]
659 #[case(1000, 200)]
660 #[case(2048, 512)]
661 fn test_sliced_to_primitive(#[case] size: usize, #[case] slice_start: usize) {
662 let mut ctx = SESSION.create_execution_ctx();
663 let values: Vec<Option<f64>> = (0..size)
664 .map(|i| {
665 if i % 5 == 0 {
666 None
667 } else if i % 4 == 3 {
668 Some(PI)
669 } else {
670 Some(1.0)
671 }
672 })
673 .collect();
674
675 let array = PrimitiveArray::from_option_iter(values.clone());
676 let encoded = alp_encode(array.as_view(), None, &mut ctx).unwrap();
677
678 let slice_end = size - slice_start;
679 let sliced_encoded = encoded.slice(slice_start..slice_end).unwrap();
680
681 let expected =
682 PrimitiveArray::from_option_iter(values[slice_start..slice_end].iter().copied());
683 assert_arrays_eq!(sliced_encoded, expected, &mut ctx);
684 }
685
686 #[test]
695 fn test_execute_decompress_with_patches_no_chunk_offsets_regression_5948() {
696 let values: Vec<f64> = vec![1.0, 2.0, PI, 4.0, 5.0];
698 let original = PrimitiveArray::from_iter(values);
699
700 let normally_encoded = alp_encode(
702 original.as_view(),
703 None,
704 &mut SESSION.create_execution_ctx(),
705 )
706 .unwrap();
707 assert!(
708 normally_encoded.patches().is_some(),
709 "Test requires patches to be present"
710 );
711
712 let original_patches = normally_encoded.patches().unwrap();
713 assert!(
714 original_patches.chunk_offsets().is_some(),
715 "Normal encoding should have chunk_offsets"
716 );
717
718 let patches_without_chunk_offsets = Patches::new(
720 original_patches.array_len(),
721 original_patches.offset(),
722 original_patches.indices().clone(),
723 original_patches.values().clone(),
724 None, )
726 .unwrap();
727
728 let alp_without_chunk_offsets = ALP::new(
730 normally_encoded.encoded().clone(),
731 normally_encoded.exponents(),
732 Some(patches_without_chunk_offsets),
733 );
734
735 let result_legacy = decompress_into_array(
737 alp_without_chunk_offsets.clone(),
738 &mut SESSION.create_execution_ctx(),
739 )
740 .unwrap();
741 let legacy_slice = result_legacy.as_slice::<f64>();
742
743 assert!(
745 (legacy_slice[2] - PI).abs() < 1e-10,
746 "Legacy path should have PI at index 2, got {}",
747 legacy_slice[2]
748 );
749
750 let result_execute = {
752 let mut ctx = SESSION.create_execution_ctx();
753 execute_decompress(alp_without_chunk_offsets, &mut ctx).unwrap()
754 };
755 let execute_slice = result_execute.as_slice::<f64>();
756
757 assert!(
759 (execute_slice[2] - PI).abs() < 1e-10,
760 "Execute path should have PI at index 2, but got {} (patches were dropped!)",
761 execute_slice[2]
762 );
763 }
764}