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

vortex_alp/alp/
decompress.rs
