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