fix: ALP-RD scalar_at respects nullability (#1864)

spiraldb · Jan 9, 2025 · 5696233 · 5696233
1 parent 8709041
commit 5696233
Showing 1 changed file with 32 additions and 2 deletions.
diff --git a/encodings/alp/src/alp_rd/compute/scalar_at.rs b/encodings/alp/src/alp_rd/compute/scalar_at.rs
@@ -1,4 +1,5 @@
 use vortex_array::compute::{scalar_at, ScalarAtFn};
+use vortex_array::ArrayDType;
 use vortex_error::VortexResult;
 use vortex_scalar::Scalar;
 
@@ -26,11 +27,11 @@ impl ScalarAtFn<ALPRDArray> for ALPRDEncoding {
         if array.is_f32() {
             let right: u32 = scalar_at(array.right_parts(), index)?.try_into()?;
             let packed = f32::from_bits((left as u32) << array.right_bit_width() | right);
-            Ok(packed.into())
+            Ok(Scalar::primitive(packed, array.dtype().nullability()))
         } else {
             let right: u64 = scalar_at(array.right_parts(), index)?.try_into()?;
             let packed = f64::from_bits(((left as u64) << array.right_bit_width()) | right);
-            Ok(packed.into())
+            Ok(Scalar::primitive(packed, array.dtype().nullability()))
         }
     }
 }
@@ -40,6 +41,7 @@ mod test {
     use rstest::rstest;
     use vortex_array::array::PrimitiveArray;
     use vortex_array::compute::scalar_at;
+    use vortex_dtype::Nullability;
     use vortex_scalar::Scalar;
 
     use crate::{ALPRDFloat, RDEncoder};
@@ -65,4 +67,32 @@ mod test {
         // The right value hits the left_part_exceptions
         assert_eq!(scalar_at(encoded.as_ref(), 2).unwrap(), outlier.into());
     }
+
+    #[test]
+    fn nullable_scalar_at() {
+        let a = 0.1f64;
+        let b = 0.2f64;
+        let outlier = 3e100f64;
+        let array = PrimitiveArray::from_option_iter([Some(a), Some(b), Some(outlier)]);
+        let encoded = RDEncoder::new(&[a, b]).encode(&array);
+
+        // Make sure that we're testing the exception pathway.
+        assert!(encoded.left_parts_patches().is_some());
+
+        // The first two values need no patching
+        assert_eq!(
+            scalar_at(encoded.as_ref(), 0).unwrap(),
+            Scalar::primitive(a, Nullability::Nullable)
+        );
+        assert_eq!(
+            scalar_at(encoded.as_ref(), 1).unwrap(),
+            Scalar::primitive(b, Nullability::Nullable)
+        );
+
+        // The right value hits the left_part_exceptions
+        assert_eq!(
+            scalar_at(encoded.as_ref(), 2).unwrap(),
+            Scalar::primitive(outlier, Nullability::Nullable)
+        );
+    }
 }