vortex_fastlanes/for/array/
for_compress.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use num_traits::PrimInt;
5use num_traits::WrappingSub;
6use vortex_array::ExecutionCtx;
7use vortex_array::IntoArray;
8use vortex_array::arrays::PrimitiveArray;
9use vortex_array::dtype::NativePType;
10use vortex_array::expr::stats::Stat;
11use vortex_array::match_each_integer_ptype;
12use vortex_error::VortexResult;
13use vortex_error::vortex_err;
14
15use crate::FoR;
16use crate::FoRArray;
17use crate::FoRData;
18
19impl FoRData {
20    pub fn encode(array: PrimitiveArray, ctx: &mut ExecutionCtx) -> VortexResult<FoRArray> {
21        let array_ref = array.clone().into_array();
22        let min = array_ref
23            .statistics()
24            .compute_stat(Stat::Min, ctx)?
25            .ok_or_else(|| vortex_err!("Min stat not found"))?;
26
27        let encoded = match_each_integer_ptype!(array.ptype(), |T| {
28            compress_primitive::<T>(array, T::try_from(&min)?, ctx)?.into_array()
29        });
30        FoR::try_new(encoded, min)
31    }
32}
33
34fn compress_primitive<T: NativePType + WrappingSub + PrimInt>(
35    parray: PrimitiveArray,
36    min: T,
37    ctx: &mut ExecutionCtx,
38) -> VortexResult<PrimitiveArray> {
39    // Set null values to the min value, ensuring that decompress into a value in the primitive
40    // range (and stop them wrapping around).
41    let encoded = parray.map_each_with_validity::<T, _, _>(ctx, |(v, bool)| {
42        if bool {
43            v.wrapping_sub(&min)
44        } else {
45            T::zero()
46        }
47    })?;
48    Ok(encoded)
49}
50
51#[cfg(test)]
52mod test {
53    use std::sync::LazyLock;
54
55    use itertools::Itertools;
56    use vortex_array::VortexSessionExecute;
57    use vortex_array::array_session;
58    use vortex_array::arrays::primitive::PrimitiveArrayExt;
59    use vortex_array::assert_arrays_eq;
60    use vortex_array::dtype::PType;
61    use vortex_array::expr::stats::StatsProvider;
62    use vortex_array::scalar::Scalar;
63    use vortex_array::validity::Validity;
64    use vortex_buffer::Buffer;
65    use vortex_buffer::buffer;
66    use vortex_session::VortexSession;
67
68    use super::*;
69    use crate::BitPackedData;
70    use crate::r#for::array::FoRArrayExt;
71    use crate::r#for::array::FoRArraySlotsExt;
72    use crate::r#for::array::for_decompress::decompress;
73    use crate::r#for::array::for_decompress::fused_decompress;
74
75    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
76        let session = array_session();
77        crate::initialize(&session);
78        session
79    });
80
81    #[test]
82    fn test_compress_round_trip_small() {
83        let mut ctx = SESSION.create_execution_ctx();
84        let array = PrimitiveArray::new((1i32..10).collect::<Buffer<_>>(), Validity::NonNullable);
85        let compressed = FoRData::encode(array.clone(), &mut ctx).unwrap();
86        assert_eq!(i32::try_from(compressed.reference_scalar()).unwrap(), 1);
87
88        assert_arrays_eq!(compressed, array, &mut ctx);
89    }
90
91    #[test]
92    fn test_compress() {
93        let mut ctx = SESSION.create_execution_ctx();
94        // Create a range offset by a million.
95        let array = PrimitiveArray::new(
96            (0u32..10_000).map(|v| v + 1_000_000).collect::<Buffer<_>>(),
97            Validity::NonNullable,
98        );
99        let compressed = FoRData::encode(array, &mut ctx).unwrap();
100        assert_eq!(
101            u32::try_from(compressed.reference_scalar()).unwrap(),
102            1_000_000u32
103        );
104    }
105
106    #[test]
107    fn test_zeros() {
108        let mut ctx = SESSION.create_execution_ctx();
109        let array = PrimitiveArray::new(buffer![0i32; 100], Validity::NonNullable);
110        assert_eq!(array.statistics().len(), 0);
111
112        let dtype = array.dtype().clone();
113        let compressed = FoRData::encode(array, &mut ctx).unwrap();
114        assert_eq!(compressed.reference_scalar().dtype(), &dtype);
115        assert!(compressed.reference_scalar().dtype().is_signed_int());
116        assert!(compressed.encoded().dtype().is_signed_int());
117
118        let encoded = compressed.encoded().execute_scalar(0, &mut ctx).unwrap();
119        assert_eq!(encoded, Scalar::from(0i32));
120    }
121
122    #[test]
123    fn test_decompress() {
124        let mut ctx = SESSION.create_execution_ctx();
125        // Create a range offset by a million.
126        let array = PrimitiveArray::from_iter((0u32..100_000).step_by(1024).map(|v| v + 1_000_000));
127        let compressed = FoRData::encode(array.clone(), &mut ctx).unwrap();
128        assert_arrays_eq!(compressed, array, &mut ctx);
129    }
130
131    #[test]
132    fn test_decompress_fused() {
133        let mut ctx = SESSION.create_execution_ctx();
134        // Create a range offset by a million.
135        let expect = PrimitiveArray::from_iter((0u32..1024).map(|x| x % 7 + 10));
136        let array = PrimitiveArray::from_iter((0u32..1024).map(|x| x % 7));
137        let bp = BitPackedData::encode(&array.into_array(), 3, &mut ctx).unwrap();
138        let compressed = FoR::try_new(bp.into_array(), 10u32.into()).unwrap();
139        assert_arrays_eq!(compressed, expect, &mut ctx);
140    }
141
142    #[test]
143    fn test_decompress_fused_patches() -> VortexResult<()> {
144        let mut ctx = SESSION.create_execution_ctx();
145        // Create a range offset by a million.
146        let expect = PrimitiveArray::from_iter((0u32..1024).map(|x| x % 7 + 10));
147        let array = PrimitiveArray::from_iter((0u32..1024).map(|x| x % 7));
148        let bp = BitPackedData::encode(&array.into_array(), 2, &mut ctx)?;
149        let compressed = FoR::try_new(bp.clone().into_array(), 10u32.into())?;
150        let decompressed = fused_decompress::<u32>(&compressed, bp.as_view(), &mut ctx)?;
151        assert_arrays_eq!(decompressed, expect, &mut ctx);
152        Ok(())
153    }
154
155    #[test]
156    fn test_overflow() -> VortexResult<()> {
157        let mut ctx = SESSION.create_execution_ctx();
158        let array = PrimitiveArray::from_iter(i8::MIN..=i8::MAX);
159        let compressed = FoRData::encode(array.clone(), &mut ctx)?;
160        assert_eq!(
161            i8::MIN,
162            compressed
163                .reference_scalar()
164                .as_primitive()
165                .typed_value::<i8>()
166                .unwrap()
167        );
168
169        let encoded = compressed
170            .encoded()
171            .clone()
172            .execute::<PrimitiveArray>(&mut ctx)?
173            .reinterpret_cast(PType::U8);
174        let unsigned: Vec<u8> = (0..=u8::MAX).collect_vec();
175        let expected_unsigned = PrimitiveArray::from_iter(unsigned);
176        assert_eq!(encoded.as_slice::<u8>(), expected_unsigned.as_slice::<u8>());
177
178        let decompressed = decompress(&compressed, &mut ctx)?;
179        array
180            .as_slice::<i8>()
181            .iter()
182            .enumerate()
183            .for_each(|(i, v)| {
184                assert_eq!(
185                    *v,
186                    i8::try_from(&compressed.execute_scalar(i, &mut ctx).unwrap()).unwrap()
187                );
188            });
189        assert_arrays_eq!(decompressed, array, &mut ctx);
190        Ok(())
191    }
192}
vortex_fastlanes/for/array/for_compress.rs

vortex_fastlanes/for/array/
for_compress.rs
