Zigzag encode/decode reuses the underlying data vec (#1638)

spiraldb · Dec 10, 2024 · 8c8dac0 · 8c8dac0
1 parent 02b1ef3
commit 8c8dac0
Showing 1 changed file with 13 additions and 16 deletions.
diff --git a/encodings/zigzag/src/compress.rs b/encodings/zigzag/src/compress.rs
@@ -1,5 +1,4 @@
 use vortex_array::array::PrimitiveArray;
-use vortex_array::stats::ArrayStatistics as _;
 use vortex_array::validity::Validity;
 use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::IntoArrayData;
@@ -12,37 +11,35 @@ use crate::ZigZagArray;
 pub fn zigzag_encode(parray: PrimitiveArray) -> VortexResult<ZigZagArray> {
     let validity = parray.validity();
     let encoded = match parray.ptype() {
-        PType::I8 => zigzag_encode_primitive::<i8>(parray.maybe_null_slice(), validity),
-        PType::I16 => zigzag_encode_primitive::<i16>(parray.maybe_null_slice(), validity),
-        PType::I32 => zigzag_encode_primitive::<i32>(parray.maybe_null_slice(), validity),
-        PType::I64 => zigzag_encode_primitive::<i64>(parray.maybe_null_slice(), validity),
+        PType::I8 => zigzag_encode_primitive::<i8>(parray.into_maybe_null_slice(), validity),
+        PType::I16 => zigzag_encode_primitive::<i16>(parray.into_maybe_null_slice(), validity),
+        PType::I32 => zigzag_encode_primitive::<i32>(parray.into_maybe_null_slice(), validity),
+        PType::I64 => zigzag_encode_primitive::<i64>(parray.into_maybe_null_slice(), validity),
         _ => vortex_bail!(
             "ZigZag can only encode signed integers, got {}",
             parray.ptype()
         ),
     };
-    let zz = ZigZagArray::try_new(encoded.into_array())?;
-    zz.inherit_statistics(parray.statistics());
-    Ok(zz)
+    ZigZagArray::try_new(encoded.into_array())
 }
 
 fn zigzag_encode_primitive<T: ExternalZigZag + NativePType>(
-    values: &[T],
+    values: Vec<T>,
     validity: Validity,
 ) -> PrimitiveArray
 where
     <T as ExternalZigZag>::UInt: NativePType,
 {
-    PrimitiveArray::from_vec(values.iter().map(|v| T::encode(*v)).collect(), validity)
+    PrimitiveArray::from_vec(values.into_iter().map(T::encode).collect(), validity)
 }
 
 pub fn zigzag_decode(parray: PrimitiveArray) -> VortexResult<PrimitiveArray> {
     let validity = parray.validity();
     let decoded = match parray.ptype() {
-        PType::U8 => zigzag_decode_primitive::<i8>(parray.maybe_null_slice(), validity),
-        PType::U16 => zigzag_decode_primitive::<i16>(parray.maybe_null_slice(), validity),
-        PType::U32 => zigzag_decode_primitive::<i32>(parray.maybe_null_slice(), validity),
-        PType::U64 => zigzag_decode_primitive::<i64>(parray.maybe_null_slice(), validity),
+        PType::U8 => zigzag_decode_primitive::<i8>(parray.into_maybe_null_slice(), validity),
+        PType::U16 => zigzag_decode_primitive::<i16>(parray.into_maybe_null_slice(), validity),
+        PType::U32 => zigzag_decode_primitive::<i32>(parray.into_maybe_null_slice(), validity),
+        PType::U64 => zigzag_decode_primitive::<i64>(parray.into_maybe_null_slice(), validity),
         _ => vortex_bail!(
             "ZigZag can only decode unsigned integers, got {}",
             parray.ptype()
@@ -52,13 +49,13 @@ pub fn zigzag_decode(parray: PrimitiveArray) -> VortexResult<PrimitiveArray> {
 }
 
 fn zigzag_decode_primitive<T: ExternalZigZag + NativePType>(
-    values: &[T::UInt],
+    values: Vec<T::UInt>,
     validity: Validity,
 ) -> PrimitiveArray
 where
     <T as ExternalZigZag>::UInt: NativePType,
 {
-    PrimitiveArray::from_vec(values.iter().map(|v| T::decode(*v)).collect(), validity)
+    PrimitiveArray::from_vec(values.into_iter().map(T::decode).collect(), validity)
 }
 
 #[cfg(test)]