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

vortex_alp/alp/
decompress.rs
