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 vortex_array::IntoArray;
81    use vortex_array::LEGACY_SESSION;
82    use vortex_array::VortexSessionExecute;
83    use vortex_array::assert_arrays_eq;
84
85    use super::*;
86    use crate::ZigZag;
87
88    #[test]
89    fn test_compress_i8() {
90        let mut ctx = LEGACY_SESSION.create_execution_ctx();
91        let compressed = zigzag_encode(PrimitiveArray::from_iter(-100_i8..100).as_view())
92            .unwrap()
93            .into_array();
94        assert!(compressed.is::<ZigZag>());
95        let decompressed = compressed.execute::<PrimitiveArray>(&mut ctx).unwrap();
96        assert_arrays_eq!(decompressed, PrimitiveArray::from_iter(-100_i8..100));
97    }
98    #[test]
99    fn test_compress_i16() {
100        let mut ctx = LEGACY_SESSION.create_execution_ctx();
101        let compressed = zigzag_encode(PrimitiveArray::from_iter(-100_i16..100).as_view())
102            .unwrap()
103            .into_array();
104        assert!(compressed.is::<ZigZag>());
105        let decompressed = compressed.execute::<PrimitiveArray>(&mut ctx).unwrap();
106        assert_arrays_eq!(decompressed, PrimitiveArray::from_iter(-100_i16..100));
107    }
108    #[test]
109    fn test_compress_i32() {
110        let mut ctx = LEGACY_SESSION.create_execution_ctx();
111        let compressed = zigzag_encode(PrimitiveArray::from_iter(-100_i32..100).as_view())
112            .unwrap()
113            .into_array();
114        assert!(compressed.is::<ZigZag>());
115        let decompressed = compressed.execute::<PrimitiveArray>(&mut ctx).unwrap();
116        assert_arrays_eq!(decompressed, PrimitiveArray::from_iter(-100_i32..100));
117    }
118    #[test]
119    fn test_compress_i64() {
120        let mut ctx = LEGACY_SESSION.create_execution_ctx();
121        let compressed = zigzag_encode(PrimitiveArray::from_iter(-100_i64..100).as_view())
122            .unwrap()
123            .into_array();
124        assert!(compressed.is::<ZigZag>());
125        let decompressed = compressed.execute::<PrimitiveArray>(&mut ctx).unwrap();
126        assert_arrays_eq!(decompressed, PrimitiveArray::from_iter(-100_i64..100));
127    }
128}
vortex_zigzag/compress.rs

vortex_zigzag/
compress.rs
