Fix alp null handling (#623)

Closes #621
spiraldb · Aug 14, 2024 · 9414b70 · 9414b70
1 parent 79a22c3
commit 9414b70
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 3 deletions.
diff --git a/encodings/alp/src/compress.rs b/encodings/alp/src/compress.rs
@@ -106,6 +106,11 @@ fn decompress_primitive<T: NativePType + ALPFloat>(
 
 #[cfg(test)]
 mod tests {
+    use core::f64;
+
+    use vortex::compute::unary::scalar_at;
+    use vortex::AsArray;
+
     use super::*;
 
     #[test]
@@ -143,7 +148,7 @@ mod tests {
     }
 
     #[test]
-    #[allow(clippy::approx_constant)]
+    #[allow(clippy::approx_constant)] // ALP doesn't like E
     fn test_patched_compress() {
         let values = vec![1.234f64, 2.718, std::f64::consts::PI, 4.0];
         let array = PrimitiveArray::from(values.clone());
@@ -158,4 +163,31 @@ mod tests {
         let decoded = decompress(encoded).unwrap();
         assert_eq!(values, decoded.maybe_null_slice::<f64>());
     }
+
+    #[test]
+    #[allow(clippy::approx_constant)] // ALP doesn't like E
+    fn test_nullable_patched_scalar_at() {
+        let values = vec![
+            Some(1.234f64),
+            Some(2.718),
+            Some(std::f64::consts::PI),
+            Some(4.0),
+            None,
+        ];
+        let array = PrimitiveArray::from_nullable_vec(values);
+        let encoded = alp_encode(&array).unwrap();
+        assert!(encoded.patches().is_some());
+
+        assert_eq!(encoded.exponents(), Exponents { e: 3, f: 0 });
+
+        for idx in 0..3 {
+            let s = scalar_at(encoded.as_array_ref(), idx).unwrap();
+            assert!(s.is_valid());
+        }
+
+        let s = scalar_at(encoded.as_array_ref(), 4).unwrap();
+        assert!(s.is_null());
+
+        let _decoded = decompress(encoded).unwrap();
+    }
 }
diff --git a/encodings/alp/src/compute.rs b/encodings/alp/src/compute.rs
@@ -1,6 +1,7 @@
 use vortex::compute::unary::{scalar_at, ScalarAtFn};
 use vortex::compute::{slice, take, ArrayCompute, SliceFn, TakeFn};
 use vortex::{Array, IntoArray};
+use vortex_dtype::{DType, NativePType};
 use vortex_error::VortexResult;
 use vortex_scalar::Scalar;
 
@@ -22,11 +23,23 @@ impl ArrayCompute for ALPArray {
 
 impl ScalarAtFn for ALPArray {
     fn scalar_at(&self, index: usize) -> VortexResult<Scalar> {
-        if let Some(patch) = self.patches().and_then(|p| scalar_at(&p, index).ok()) {
-            return Ok(patch);
+        if let Some(patches) = self.patches().and_then(|p| {
+            p.with_dyn(|arr| {
+                // We need to make sure the value is actually in the patches array
+                arr.is_valid(index)
+            })
+            .then_some(p)
+        }) {
+            return scalar_at(&patches, index);
         }
+
         let encoded_val = scalar_at(&self.encoded(), index)?;
+
         match_each_alp_float_ptype!(self.ptype(), |$T| {
+            // If we have a null value, no need to decode it
+            if !encoded_val.is_valid() {
+                return Ok(Scalar::null(DType::from(<$T>::PTYPE).as_nullable()));
+            }
             let encoded_val: <$T as ALPFloat>::ALPInt = encoded_val.as_ref().try_into().unwrap();
             Ok(Scalar::from(<$T as ALPFloat>::decode_single(
                 encoded_val,