vortex_zigzag/
compress.rs

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

vortex_zigzag/
compress.rs
