vortex_zigzag/
compress.rs

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

vortex_zigzag/
compress.rs
