vortex_zigzag/
compress.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::arrays::PrimitiveArray;
5use vortex_array::validity::Validity;
6use vortex_array::vtable::ValidityHelper;
7use vortex_buffer::BufferMut;
8use vortex_dtype::{NativePType, PType};
9use vortex_error::{VortexResult, vortex_bail, vortex_panic};
10use zigzag::ZigZag as ExternalZigZag;
11
12use crate::ZigZagArray;
13
14pub fn zigzag_encode(parray: PrimitiveArray) -> VortexResult<ZigZagArray> {
15    let validity = parray.validity().clone();
16    let encoded = match parray.ptype() {
17        PType::I8 => zigzag_encode_primitive::<i8>(parray.into_buffer_mut(), validity),
18        PType::I16 => zigzag_encode_primitive::<i16>(parray.into_buffer_mut(), validity),
19        PType::I32 => zigzag_encode_primitive::<i32>(parray.into_buffer_mut(), validity),
20        PType::I64 => zigzag_encode_primitive::<i64>(parray.into_buffer_mut(), validity),
21        _ => vortex_bail!(
22            "ZigZag can only encode signed integers, got {}",
23            parray.ptype()
24        ),
25    };
26    ZigZagArray::try_new(encoded.to_array())
27}
28
29fn zigzag_encode_primitive<T: ExternalZigZag + NativePType>(
30    values: BufferMut<T>,
31    validity: Validity,
32) -> PrimitiveArray
33where
34    <T as ExternalZigZag>::UInt: NativePType,
35{
36    PrimitiveArray::new(
37        values.map_each_in_place(|v| T::encode(v)).freeze(),
38        validity,
39    )
40}
41
42pub fn zigzag_decode(parray: PrimitiveArray) -> PrimitiveArray {
43    let validity = parray.validity().clone();
44    match parray.ptype() {
45        PType::U8 => zigzag_decode_primitive::<i8>(parray.into_buffer_mut(), validity),
46        PType::U16 => zigzag_decode_primitive::<i16>(parray.into_buffer_mut(), validity),
47        PType::U32 => zigzag_decode_primitive::<i32>(parray.into_buffer_mut(), validity),
48        PType::U64 => zigzag_decode_primitive::<i64>(parray.into_buffer_mut(), validity),
49        _ => vortex_panic!(
50            "ZigZag can only decode unsigned integers, got {}",
51            parray.ptype()
52        ),
53    }
54}
55
56fn zigzag_decode_primitive<T: ExternalZigZag + NativePType>(
57    values: BufferMut<T::UInt>,
58    validity: Validity,
59) -> PrimitiveArray
60where
61    <T as ExternalZigZag>::UInt: NativePType,
62{
63    PrimitiveArray::new(
64        values.map_each_in_place(|v| T::decode(v)).freeze(),
65        validity,
66    )
67}
68
69#[cfg(test)]
70mod test {
71    use vortex_array::ToCanonical;
72
73    use super::*;
74    use crate::ZigZagEncoding;
75
76    #[test]
77    fn test_compress_i8() {
78        let compressed = zigzag_encode(PrimitiveArray::from_iter(-100_i8..100)).unwrap();
79        assert_eq!(compressed.encoding_id(), ZigZagEncoding.id());
80        assert_eq!(
81            compressed.to_primitive().as_slice::<i8>(),
82            (-100_i8..100).collect::<Vec<_>>()
83        );
84    }
85    #[test]
86    fn test_compress_i16() {
87        let compressed = zigzag_encode(PrimitiveArray::from_iter(-100_i16..100)).unwrap();
88        assert_eq!(compressed.encoding_id(), ZigZagEncoding.id());
89        assert_eq!(
90            compressed.to_primitive().as_slice::<i16>(),
91            (-100_i16..100).collect::<Vec<_>>()
92        );
93    }
94    #[test]
95    fn test_compress_i32() {
96        let compressed = zigzag_encode(PrimitiveArray::from_iter(-100_i32..100)).unwrap();
97        assert_eq!(compressed.encoding_id(), ZigZagEncoding.id());
98        assert_eq!(
99            compressed.to_primitive().as_slice::<i32>(),
100            (-100_i32..100).collect::<Vec<_>>()
101        );
102    }
103    #[test]
104    fn test_compress_i64() {
105        let compressed = zigzag_encode(PrimitiveArray::from_iter(-100_i64..100)).unwrap();
106        assert_eq!(compressed.encoding_id(), ZigZagEncoding.id());
107        assert_eq!(
108            compressed.to_primitive().as_slice::<i64>(),
109            (-100_i64..100).collect::<Vec<_>>()
110        );
111    }
112}
vortex_zigzag/compress.rs

vortex_zigzag/
compress.rs
