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        // SAFETY: the exception positions are collected in increasing order from positions in
108        // the input, with one exception value each.
109        unsafe {
110            Patches::new_unchecked(
111                len,
112                0,
113                packed_pos,
114                Buffer::from(exceptions).into_array(),
115                // TODO(0ax1): handle chunk offsets
116                None,
117                None,
118            )
119        }
120    });
121
122    ALPRD::try_new(
123        DType::Primitive(T::PTYPE, packed_left.dtype().nullability()),
124        packed_left,
125        Buffer::<u16>::copy_from(encoder.codes()),
126        packed_right,
127        encoder.right_bit_width(),
128        exceptions,
129    )
130    .vortex_expect("ALPRDArray construction in encode")
131}
132
133/// Decode ALP-RD encoded values back into their original floating point format.
134///
135/// The decoded values are written over `right_parts`, which shares its layout with the floats they
136/// decode to.
137///
138/// # Panics
139///
140/// Panics if `left_parts` and `right_parts` differ in length.
141pub fn alp_rd_decode<T: ALPRDFloat + NativePType>(
142    mut left_parts: BufferMut<u16>,
143    left_parts_dict: &[u16],
144    right_bit_width: u8,
145    mut right_parts: BufferMut<T::UINT>,
146    left_parts_patches: Option<Patches>,
147    ctx: &mut ExecutionCtx,
148) -> VortexResult<Buffer<T>>
149where
150    T::UINT: NativePType,
151{
152    if left_parts.len() != right_parts.len() {
153        vortex_panic!("alp_rd_decode: left_parts.len != right_parts.len");
154    }
155
156    if let Some(patches) = left_parts_patches {
157        // Patched path: some left-part codes map to exception values that live outside
158        // the dictionary. We must dictionary-decode first, then overwrite the exceptions,
159        // before we can combine with right-parts.
160        alp_rd_dict_decode_inplace(left_parts.as_mut_slice(), left_parts_dict);
161
162        let indices = patches.indices().clone().execute::<PrimitiveArray>(ctx)?;
163        let patch_values = patches.values().clone().execute::<PrimitiveArray>(ctx)?;
164        match_each_integer_ptype!(indices.ptype(), |I| {
165            alp_rd_apply_patches(
166                left_parts.as_mut_slice(),
167                indices.as_slice::<I>(),
168                patch_values.as_slice::<u16>(),
169                patches.offset(),
170            )
171        });
172
173        alp_rd_combine_inplace::<T>(
174            right_parts.as_mut_slice(),
175            left_parts.as_ref(),
176            right_bit_width,
177        );
178    } else {
179        // Non-patched fast path: every code maps through the dictionary, so the hot loop is a
180        // lookup into the pre-shifted dictionary and an OR.
181        alp_rd_combine_codes_inplace::<T>(
182            right_parts.as_mut_slice(),
183            left_parts.as_ref(),
184            left_parts_dict,
185            right_bit_width,
186        );
187    }
188
189    // SAFETY: all bit patterns of T::UINT are valid T (u32↔f32 or u64↔f64).
190    Ok(unsafe { right_parts.transmute::<T>() }.freeze())
191}
192
193#[cfg(test)]
194mod tests {
195    use vortex_array::arrays::PrimitiveArray;
196
197    use super::ALPRDArrayExt;
198    use super::RDEncoder;
199    use super::RDEncoderExt;
200
201    #[test]
202    fn rd_encoder_uses_deterministic_codes_for_tied_prefixes() {
203        let values: Vec<_> = (1_000_i64..2_000)
204            .map(|row_id| ((row_id * 53) % 36_000) as f64 / 100.0 - 180.0)
205            .collect();
206        let array = PrimitiveArray::from_iter(values.iter().copied());
207
208        let encoded = RDEncoder::new(&values).encode(array.as_view());
209
210        assert_eq!(encoded.right_bit_width(), 53);
211        assert_eq!(
212            encoded.left_parts_dictionary().as_slice(),
213            &[514, 1538, 515, 1539, 513, 1537, 512, 1536]
214        );
215    }
216}
vortex_alp/alp_rd/mod.rs

vortex_alp/alp_rd/
mod.rs
