diff --git a/encodings/alp/src/alp_rd/compute/scalar_at.rs b/encodings/alp/src/alp_rd/compute/scalar_at.rs index d8c0c49e2e..c026ca2086 100644 --- 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 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) + ); + } }