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

vortex_zigzag/
compress.rs
