vortex_zigzag/
compress.rs

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

vortex_zigzag/
compress.rs
