vortex_alp/alp/
decompress.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::mem::transmute;
5
6use vortex_array::ExecutionCtx;
7use vortex_array::ToCanonical;
8use vortex_array::arrays::PrimitiveArray;
9use vortex_array::arrays::chunk_range;
10use vortex_array::arrays::patch_chunk;
11use vortex_array::dtype::DType;
12use vortex_array::match_each_unsigned_integer_ptype;
13use vortex_array::patches::Patches;
14use vortex_array::vtable::ValidityHelper;
15use vortex_buffer::BufferMut;
16use vortex_error::VortexResult;
17
18use crate::ALPArray;
19use crate::ALPFloat;
20use crate::Exponents;
21use crate::match_each_alp_float_ptype;
22
23/// Decompresses an ALP-encoded array using `to_primitive` (legacy path).
24///
25/// # Returns
26///
27/// A `PrimitiveArray` containing the decompressed floating-point values with all patches applied.
28pub fn decompress_into_array(
29    array: ALPArray,
30    ctx: &mut ExecutionCtx,
31) -> VortexResult<PrimitiveArray> {
32    let (encoded, exponents, patches, dtype) = array.into_parts();
33    if let Some(ref patches) = patches
34        && let Some(chunk_offsets) = patches.chunk_offsets()
35    {
36        let prim_encoded = encoded.to_primitive();
37        // We need to drop ALPArray here in case converting encoded buffer into
38        // primitive didn't create a copy. In that case both alp_encoded and array
39        // will hold a reference to the buffer we want to mutate.
40        drop(encoded);
41        let patches_chunk_offsets = chunk_offsets.as_ref().to_primitive();
42        let patches_indices = patches.indices().to_primitive();
43        let patches_values = patches.values().to_primitive();
44        Ok(decompress_chunked_core(
45            prim_encoded,
46            exponents,
47            &patches_indices,
48            &patches_values,
49            &patches_chunk_offsets,
50            patches,
51            dtype,
52        ))
53    } else {
54        let encoded_prim = encoded.to_primitive();
55        // We need to drop ALPArray here in case converting encoded buffer into
56        // primitive didn't create a copy. In that case both alp_encoded and array
57        // will hold a reference to the buffer we want to mutate.
58        drop(encoded);
59        decompress_unchunked_core(encoded_prim, exponents, patches, dtype, ctx)
60    }
61}
62
63/// Decompresses an ALP-encoded array using `execute` (execution path).
64///
65/// This version uses `execute` on child arrays instead of `to_primitive`,
66/// ensuring proper recursive execution through the execution context.
67///
68/// # Returns
69///
70/// A `PrimitiveArray` containing the decompressed floating-point values with all patches applied.
71pub fn execute_decompress(array: ALPArray, ctx: &mut ExecutionCtx) -> VortexResult<PrimitiveArray> {
72    let (encoded, exponents, patches, dtype) = array.into_parts();
73    if let Some(ref patches) = patches
74        && let Some(chunk_offsets) = patches.chunk_offsets()
75    {
76        // TODO(joe): have into parts.
77        let encoded = encoded.execute::<PrimitiveArray>(ctx)?;
78        let patches_chunk_offsets = chunk_offsets.clone().execute::<PrimitiveArray>(ctx)?;
79        let patches_indices = patches.indices().clone().execute::<PrimitiveArray>(ctx)?;
80        let patches_values = patches.values().clone().execute::<PrimitiveArray>(ctx)?;
81        Ok(decompress_chunked_core(
82            encoded,
83            exponents,
84            &patches_indices,
85            &patches_values,
86            &patches_chunk_offsets,
87            patches,
88            dtype,
89        ))
90    } else {
91        let encoded = encoded.execute::<PrimitiveArray>(ctx)?;
92        decompress_unchunked_core(encoded, exponents, patches, dtype, ctx)
93    }
94}
95
96/// Core decompression logic for chunked ALP arrays.
97///
98/// Takes pre-resolved `PrimitiveArray` inputs to avoid duplication between
99/// the `to_primitive` and `execute` paths.
100#[expect(
101    clippy::cognitive_complexity,
102    reason = "complexity is from nested match_each_* macros"
103)]
104fn decompress_chunked_core(
105    encoded: PrimitiveArray,
106    exponents: Exponents,
107    patches_indices: &PrimitiveArray,
108    patches_values: &PrimitiveArray,
109    patches_chunk_offsets: &PrimitiveArray,
110    patches: &Patches,
111    dtype: DType,
112) -> PrimitiveArray {
113    let validity = encoded.validity().clone();
114    let ptype = dtype.as_ptype();
115    let array_len = encoded.len();
116    let offset_within_chunk = patches.offset_within_chunk().unwrap_or(0);
117
118    match_each_alp_float_ptype!(ptype, |T| {
119        let patches_values = patches_values.as_slice::<T>();
120        let mut alp_buffer = encoded.into_buffer_mut();
121        match_each_unsigned_integer_ptype!(patches_chunk_offsets.ptype(), |C| {
122            let patches_chunk_offsets = patches_chunk_offsets.as_slice::<C>();
123
124            match_each_unsigned_integer_ptype!(patches_indices.ptype(), |I| {
125                let patches_indices = patches_indices.as_slice::<I>();
126
127                for chunk_idx in 0..patches_chunk_offsets.len() {
128                    let chunk_range = chunk_range(chunk_idx, patches.offset(), array_len);
129                    let chunk_slice = &mut alp_buffer.as_mut_slice()[chunk_range];
130
131                    <T>::decode_slice_inplace(chunk_slice, exponents);
132
133                    let decoded_chunk: &mut [T] = unsafe { transmute(chunk_slice) };
134                    patch_chunk(
135                        decoded_chunk,
136                        patches_indices,
137                        patches_values,
138                        patches.offset(),
139                        patches_chunk_offsets,
140                        chunk_idx,
141                        offset_within_chunk,
142                    );
143                }
144
145                let decoded_buffer: BufferMut<T> = unsafe { transmute(alp_buffer) };
146                PrimitiveArray::new::<T>(decoded_buffer.freeze(), validity)
147            })
148        })
149    })
150}
151
152/// Core decompression logic for unchunked ALP arrays.
153///
154/// Takes a pre-resolved `PrimitiveArray` to avoid duplication between
155/// the `to_primitive` and `execute` paths.
156fn decompress_unchunked_core(
157    encoded: PrimitiveArray,
158    exponents: Exponents,
159    patches: Option<Patches>,
160    dtype: DType,
161    ctx: &mut ExecutionCtx,
162) -> VortexResult<PrimitiveArray> {
163    let validity = encoded.validity().clone();
164    let ptype = dtype.as_ptype();
165
166    let decoded = match_each_alp_float_ptype!(ptype, |T| {
167        let mut alp_buffer = encoded.into_buffer_mut();
168        <T>::decode_slice_inplace(alp_buffer.as_mut_slice(), exponents);
169        let decoded_buffer: BufferMut<T> = unsafe { transmute(alp_buffer) };
170        PrimitiveArray::new::<T>(decoded_buffer.freeze(), validity)
171    });
172
173    if let Some(patches) = patches {
174        decoded.patch(&patches, ctx)
175    } else {
176        Ok(decoded)
177    }
178}
vortex_alp/alp/decompress.rs

vortex_alp/alp/
decompress.rs
