vortex_zigzag/
compress.rs

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

vortex_zigzag/
compress.rs
