fix delta compressor for signed ints

spiraldb · Sep 23, 2024 · d4448e8 · d4448e8
1 parent f3344c2
commit d4448e8
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 10 deletions.
diff --git a/encodings/fastlanes/src/delta/compress.rs b/encodings/fastlanes/src/delta/compress.rs
@@ -3,14 +3,26 @@ use fastlanes::{Delta, Transpose};
 use num_traits::{WrappingAdd, WrappingSub};
 use vortex::array::PrimitiveArray;
 use vortex::compute::unary::fill_forward;
+use vortex::stats::ArrayStatistics;
 use vortex::validity::Validity;
 use vortex::IntoArrayVariant;
-use vortex_dtype::{match_each_unsigned_integer_ptype, NativePType, Nullability};
+use vortex_dtype::{
+    match_each_signed_integer_ptype, match_each_unsigned_integer_ptype, NativePType, Nullability,
+};
 use vortex_error::VortexResult;
 
 use crate::DeltaArray;
 
 pub fn delta_compress(array: &PrimitiveArray) -> VortexResult<(PrimitiveArray, PrimitiveArray)> {
+    let array = if array.ptype().is_signed_int() {
+        match_each_signed_integer_ptype!(array.ptype(), |$T| {
+            assert!(array.statistics().compute_min::<$T>().map(|x| x >= 0).unwrap_or(false));
+        });
+        &array.reinterpret_cast(array.ptype().to_unsigned())
+    } else {
+        array
+    };
+
     // Fill forward nulls
     let filled = fill_forward(array.as_ref())?.into_primitive()?;
 

diff --git a/encodings/fastlanes/src/delta/mod.rs b/encodings/fastlanes/src/delta/mod.rs
@@ -8,7 +8,7 @@ use vortex::validity::{ArrayValidity, LogicalValidity, Validity, ValidityMetadat
 use vortex::variants::{ArrayVariants, PrimitiveArrayTrait};
 use vortex::visitor::{AcceptArrayVisitor, ArrayVisitor};
 use vortex::{impl_encoding, Array, ArrayDType, ArrayDef, ArrayTrait, Canonical, IntoCanonical};
-use vortex_dtype::match_each_unsigned_integer_ptype;
+use vortex_dtype::{match_each_unsigned_integer_ptype, PType};
 use vortex_error::{vortex_bail, vortex_panic, VortexExpect as _, VortexResult};
 
 mod compress;
@@ -74,14 +74,14 @@ impl DeltaArray {
 
     #[inline]
     fn lanes(&self) -> usize {
-        let ptype = self.dtype().try_into().unwrap_or_else(|err| {
+        let ptype: PType = self.dtype().try_into().unwrap_or_else(|err| {
             vortex_panic!(
                 err,
                 "Failed to convert DeltaArray DType {} to PType",
                 self.dtype()
             )
         });
-        match_each_unsigned_integer_ptype!(ptype, |$T| {
+        match_each_unsigned_integer_ptype!(ptype.to_unsigned(), |$T| {
             <$T as fastlanes::FastLanes>::LANES
         })
     }

diff --git a/vortex-dtype/src/ptype.rs b/vortex-dtype/src/ptype.rs
@@ -164,6 +164,21 @@ macro_rules! match_each_unsigned_integer_ptype {
     })
 }
 
+#[macro_export]
+macro_rules! match_each_signed_integer_ptype {
+    ($self:expr, | $_:tt $enc:ident | $($body:tt)*) => ({
+        macro_rules! __with__ {( $_ $enc:ident ) => ( $($body)* )}
+        use $crate::PType;
+        match $self {
+            PType::I8 => __with__! { i8 },
+            PType::I16 => __with__! { i16 },
+            PType::I32 => __with__! { i32 },
+            PType::I64 => __with__! { i64 },
+            _ => panic!("Unsupported ptype {}", $self),
+        }
+    })
+}
+
 #[macro_export]
 macro_rules! match_each_float_ptype {
     ($self:expr, | $_:tt $enc:ident | $($body:tt)*) => ({

diff --git a/vortex-sampling-compressor/src/compressors/delta.rs b/vortex-sampling-compressor/src/compressors/delta.rs
@@ -49,12 +49,14 @@ impl EncodingCompressor for DeltaCompressor {
         let (bases, deltas) = delta_compress(&parray)?;
 
         // Recursively compress the bases and deltas
-        let bases = ctx
-            .named("bases")
-            .compress(bases.as_ref(), like.as_ref().and_then(|l| l.child(0)))?;
-        let deltas = ctx
-            .named("deltas")
-            .compress(deltas.as_ref(), like.as_ref().and_then(|l| l.child(1)))?;
+        let bases = ctx.named("bases").compress(
+            bases.reinterpret_cast(parray.ptype()).as_ref(),
+            like.as_ref().and_then(|l| l.child(0)),
+        )?;
+        let deltas = ctx.named("deltas").compress(
+            deltas.reinterpret_cast(parray.ptype()).as_ref(),
+            like.as_ref().and_then(|l| l.child(1)),
+        )?;
 
         Ok(CompressedArray::new(
             DeltaArray::try_new(bases.array, deltas.array, validity)?.into_array(),