vortex_zigzag/
compress.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::ArrayView;
5use vortex_array::IntoArray;
6use vortex_array::arrays::Primitive;
7use vortex_array::arrays::PrimitiveArray;
8use vortex_array::dtype::NativePType;
9use vortex_array::dtype::PType;
10use vortex_array::validity::Validity;
11use vortex_buffer::BufferMut;
12use vortex_error::VortexExpect;
13use vortex_error::VortexResult;
14use vortex_error::vortex_bail;
15use vortex_error::vortex_panic;
16use zigzag::ZigZag as ExternalZigZag;
17
18use crate::ZigZag;
19use crate::ZigZagArray;
20pub fn zigzag_encode(parray: ArrayView<'_, Primitive>) -> VortexResult<ZigZagArray> {
21    let parray = parray.into_owned();
22    let validity = parray.validity()?;
23    let encoded = match parray.ptype() {
24        PType::I8 => zigzag_encode_primitive::<i8>(parray.into_buffer_mut(), validity),
25        PType::I16 => zigzag_encode_primitive::<i16>(parray.into_buffer_mut(), validity),
26        PType::I32 => zigzag_encode_primitive::<i32>(parray.into_buffer_mut(), validity),
27        PType::I64 => zigzag_encode_primitive::<i64>(parray.into_buffer_mut(), validity),
28        _ => vortex_bail!(
29            "ZigZag can only encode signed integers, got {}",
30            parray.ptype()
31        ),
32    };
33    ZigZag::try_new(encoded.into_array())
34}
35
36fn zigzag_encode_primitive<T: ExternalZigZag + NativePType>(
37    values: BufferMut<T>,
38    validity: Validity,
39) -> PrimitiveArray
40where
41    <T as ExternalZigZag>::UInt: NativePType,
42{
43    PrimitiveArray::new(
44        values.map_each_in_place(|v| T::encode(v)).freeze(),
45        validity,
46    )
47}
48
49pub fn zigzag_decode(parray: PrimitiveArray) -> PrimitiveArray {
50    let validity = parray
51        .validity()
52        .vortex_expect("zigzag validity should be derivable");
53    match parray.ptype() {
54        PType::U8 => zigzag_decode_primitive::<i8>(parray.into_buffer_mut(), validity),
55        PType::U16 => zigzag_decode_primitive::<i16>(parray.into_buffer_mut(), validity),
56        PType::U32 => zigzag_decode_primitive::<i32>(parray.into_buffer_mut(), validity),
57        PType::U64 => zigzag_decode_primitive::<i64>(parray.into_buffer_mut(), validity),
58        _ => vortex_panic!(
59            "ZigZag can only decode unsigned integers, got {}",
60            parray.ptype()
61        ),
62    }
63}
64
65fn zigzag_decode_primitive<T: ExternalZigZag + NativePType>(
66    values: BufferMut<T::UInt>,
67    validity: Validity,
68) -> PrimitiveArray
69where
70    <T as ExternalZigZag>::UInt: NativePType,
71{
72    PrimitiveArray::new(
73        values.map_each_in_place(|v| T::decode(v)).freeze(),
74        validity,
75    )
76}
77
78#[cfg(test)]
79mod test {
80    use std::sync::LazyLock;
81
82    use rstest::rstest;
83    use vortex_array::IntoArray;
84    use vortex_array::VortexSessionExecute;
85    use vortex_array::assert_arrays_eq;
86    use vortex_session::VortexSession;
87
88    use super::*;
89    use crate::ZigZag;
90
91    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
92        let session = vortex_array::array_session();
93        crate::initialize(&session);
94        session
95    });
96
97    #[rstest]
98    #[case::i8(PrimitiveArray::from_iter(-100_i8..100))]
99    #[case::i16(PrimitiveArray::from_iter(-100_i16..100))]
100    #[case::i32(PrimitiveArray::from_iter(-100_i32..100))]
101    #[case::i64(PrimitiveArray::from_iter(-100_i64..100))]
102    fn test_compress(#[case] input: PrimitiveArray) {
103        let mut ctx = SESSION.create_execution_ctx();
104        let compressed = zigzag_encode(input.as_view()).unwrap().into_array();
105        assert!(compressed.is::<ZigZag>());
106        let decompressed = compressed.execute::<PrimitiveArray>(&mut ctx).unwrap();
107        assert_arrays_eq!(decompressed, input, &mut ctx);
108    }
109}
vortex_zigzag/compress.rs

vortex_zigzag/
compress.rs
