vortex_alp/alp_rd/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Vortex array support for ALP-RD ("real doubles"), on top of the `alp` crate's encoder.
5
6#![expect(clippy::cast_possible_truncation)]
7
8pub use ::alp::ALPRDFloat;
9pub use ::alp::RDEncoder;
10use ::alp::alp_rd_apply_patches;
11use ::alp::alp_rd_combine_codes_inplace;
12use ::alp::alp_rd_combine_inplace;
13use ::alp::alp_rd_dict_decode_inplace;
14use ::alp::bit_width;
15pub use array::*;
16use vortex_array::ArrayView;
17use vortex_array::ExecutionCtx;
18use vortex_array::IntoArray;
19use vortex_array::arrays::Primitive;
20use vortex_array::arrays::PrimitiveArray;
21use vortex_array::dtype::DType;
22use vortex_array::dtype::NativePType;
23use vortex_array::match_each_integer_ptype;
24use vortex_array::patches::Patches;
25use vortex_array::validity::Validity;
26use vortex_buffer::Buffer;
27use vortex_buffer::BufferMut;
28use vortex_error::VortexExpect;
29use vortex_error::VortexResult;
30use vortex_error::vortex_panic;
31use vortex_fastlanes::bitpack_compress::bitpack_encode_unchecked;
32use vortex_session::VortexSession;
33
34use crate::match_each_alp_float_ptype;
35
36mod array;
37mod compute;
38mod kernel;
39mod ops;
40mod rules;
41mod slice;
42
43pub(crate) fn initialize(session: &VortexSession) {
44    kernel::initialize(session);
45}
46
47/// Encoding of Vortex arrays with an ALP-RD [`RDEncoder`].
48pub trait RDEncoderExt {
49    /// Encode a set of floating point values with ALP-RD.
50    ///
51    /// Each value is split into a left and right component, which are compressed individually.
52    // TODO(joe): make fallible
53    fn encode(&self, array: ArrayView<'_, Primitive>) -> ALPRDArray;
54}
55
56impl RDEncoderExt for RDEncoder {
57    fn encode(&self, array: ArrayView<'_, Primitive>) -> ALPRDArray {
58        match_each_alp_float_ptype!(array.ptype(), |P| { encode_generic::<P>(self, array) })
59    }
60}
61
62fn encode_generic<T>(encoder: &RDEncoder, array: ArrayView<'_, Primitive>) -> ALPRDArray
63where
64    T: ALPRDFloat + NativePType,
65    T::UINT: NativePType,
66{
67    let (left_parts, right_parts, exceptions_pos, exceptions) =
68        encoder.split_parts(array.as_slice::<T>());
69    let len = left_parts.len();
70
71    // Bit-pack down the encoded left-parts array that have been dictionary encoded.
72    let primitive_left = PrimitiveArray::new(
73        Buffer::from(left_parts),
74        array
75            .validity()
76            .vortex_expect("ALP RD validity should be derivable"),
77    );
78    // SAFETY: by construction, all values in left_parts can be packed to left_bit_width.
79    let packed_left = unsafe {
80        bitpack_encode_unchecked(primitive_left, encoder.left_bit_width())
81            .vortex_expect("bitpack_encode_unchecked should succeed for left parts")
82            .into_array()
83    };
84
85    let primitive_right = PrimitiveArray::new(Buffer::from(right_parts), Validity::NonNullable);
86    // SAFETY: by construction, all values in right_parts are right_bit_width + leading zeros.
87    let packed_right = unsafe {
88        bitpack_encode_unchecked(primitive_right, encoder.right_bit_width())
89            .vortex_expect("bitpack_encode_unchecked should succeed for right parts")
90            .into_array()
91    };
92
93    // Patches for the left-parts that could not be dictionary encoded.
94    let exceptions = (!exceptions_pos.is_empty()).then(|| {
95        let max_exc_pos = exceptions_pos.last().copied().unwrap_or_default();
96        let bw = bit_width(max_exc_pos);
97
98        let exc_pos_array =
99            PrimitiveArray::new(Buffer::from(exceptions_pos), Validity::NonNullable);
100        // SAFETY: We calculate bw such that it is wide enough to hold the largest position index.
101        let packed_pos = unsafe {
102            bitpack_encode_unchecked(exc_pos_array, bw)
103                .vortex_expect("bitpack_encode_unchecked should succeed for exception positions")
104                .into_array()
105        };
106
107        Patches::new(
108            len,
109            0,
110            packed_pos,
111            Buffer::from(exceptions).into_array(),
112            // TODO(0ax1): handle chunk offsets
113            None,
114        )
115        .vortex_expect("Patches construction in encode")
116    });
117
118    ALPRD::try_new(
119        DType::Primitive(T::PTYPE, packed_left.dtype().nullability()),
120        packed_left,
121        Buffer::<u16>::copy_from(encoder.codes()),
122        packed_right,
123        encoder.right_bit_width(),
124        exceptions,
125    )
126    .vortex_expect("ALPRDArray construction in encode")
127}
128
129/// Decode ALP-RD encoded values back into their original floating point format.
130///
131/// The decoded values are written over `right_parts`, which shares its layout with the floats they
132/// decode to.
133///
134/// # Panics
135///
136/// Panics if `left_parts` and `right_parts` differ in length.
137pub fn alp_rd_decode<T: ALPRDFloat + NativePType>(
138    mut left_parts: BufferMut<u16>,
139    left_parts_dict: &[u16],
140    right_bit_width: u8,
141    mut right_parts: BufferMut<T::UINT>,
142    left_parts_patches: Option<Patches>,
143    ctx: &mut ExecutionCtx,
144) -> VortexResult<Buffer<T>>
145where
146    T::UINT: NativePType,
147{
148    if left_parts.len() != right_parts.len() {
149        vortex_panic!("alp_rd_decode: left_parts.len != right_parts.len");
150    }
151
152    if let Some(patches) = left_parts_patches {
153        // Patched path: some left-part codes map to exception values that live outside
154        // the dictionary. We must dictionary-decode first, then overwrite the exceptions,
155        // before we can combine with right-parts.
156        alp_rd_dict_decode_inplace(left_parts.as_mut_slice(), left_parts_dict);
157
158        let indices = patches.indices().clone().execute::<PrimitiveArray>(ctx)?;
159        let patch_values = patches.values().clone().execute::<PrimitiveArray>(ctx)?;
160        match_each_integer_ptype!(indices.ptype(), |I| {
161            alp_rd_apply_patches(
162                left_parts.as_mut_slice(),
163                indices.as_slice::<I>(),
164                patch_values.as_slice::<u16>(),
165                patches.offset(),
166            )
167        });
168
169        alp_rd_combine_inplace::<T>(
170            right_parts.as_mut_slice(),
171            left_parts.as_ref(),
172            right_bit_width,
173        );
174    } else {
175        // Non-patched fast path: every code maps through the dictionary, so the hot loop is a
176        // lookup into the pre-shifted dictionary and an OR.
177        alp_rd_combine_codes_inplace::<T>(
178            right_parts.as_mut_slice(),
179            left_parts.as_ref(),
180            left_parts_dict,
181            right_bit_width,
182        );
183    }
184
185    // SAFETY: all bit patterns of T::UINT are valid T (u32↔f32 or u64↔f64).
186    Ok(unsafe { right_parts.transmute::<T>() }.freeze())
187}
188
189#[cfg(test)]
190mod tests {
191    use vortex_array::arrays::PrimitiveArray;
192
193    use super::ALPRDArrayExt;
194    use super::RDEncoder;
195    use super::RDEncoderExt;
196
197    #[test]
198    fn rd_encoder_uses_deterministic_codes_for_tied_prefixes() {
199        let values: Vec<_> = (1_000_i64..2_000)
200            .map(|row_id| ((row_id * 53) % 36_000) as f64 / 100.0 - 180.0)
201            .collect();
202        let array = PrimitiveArray::from_iter(values.iter().copied());
203
204        let encoded = RDEncoder::new(&values).encode(array.as_view());
205
206        assert_eq!(encoded.right_bit_width(), 53);
207        assert_eq!(
208            encoded.left_parts_dictionary().as_slice(),
209            &[514, 1538, 515, 1539, 513, 1537, 512, 1536]
210        );
211    }
212}
vortex_alp/alp_rd/mod.rs

vortex_alp/alp_rd/
mod.rs
