Add checked arithmetic compute (#5096)

FLUPs:

* Arithmetic array to invert lhs/rhs for constant checks
* Special handling for vectors with many nulls (iterate only the sparse
non-null indices)

---------

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

Cargo.lock

@@ -193,7 +193,7 @@ pub trait ALPFloat: private::Sealed + Float + Display + NativePType {
     }
 
     fn decode_buffer(encoded: BufferMut<Self::ALPInt>, exponents: Exponents) -> BufferMut<Self> {
-        encoded.map_each(move |encoded| Self::decode_single(encoded, exponents))
+        encoded.map_each_in_place(move |encoded| Self::decode_single(encoded, exponents))
     }
 
     #[inline(always)]

encodings/alp/src/alp/mod.rs

@@ -311,7 +311,7 @@ pub fn alp_rd_decode<T: ALPRDFloat>(
     // Shift the left-parts and add in the right-parts.
     let mut index = 0;
     right_parts
-        .map_each(|right| {
+        .map_each_in_place(|right| {
             let left = values[index];
             index += 1;
             let left = <T as ALPRDFloat>::from_u16(left);

encodings/alp/src/alp_rd/mod.rs

@@ -53,7 +53,7 @@ pub fn decode_to_temporal(array: &DateTimePartsArray) -> TemporalArray {
     // are constant.
     let mut values: BufferMut<i64> = days_buf
         .into_buffer_mut::<i64>()
-        .map_each(|d| d * 86_400 * divisor);
+        .map_each_in_place(|d| d * 86_400 * divisor);
 
     if let Some(seconds) = array.seconds().as_constant() {
         let seconds = seconds

encodings/datetime-parts/src/canonical.rs

@@ -151,7 +151,9 @@ fn decompress_primitive<T: NativePType + WrappingAdd + PrimInt>(
     values: BufferMut<T>,
     min: T,
 ) -> Buffer<T> {
-    values.map_each(move |v| v.wrapping_add(&min)).freeze()
+    values
+        .map_each_in_place(move |v| v.wrapping_add(&min))
+        .freeze()
 }
 
 #[cfg(test)]

encodings/fastlanes/src/for/compress.rs

@@ -33,7 +33,10 @@ fn zigzag_encode_primitive<T: ExternalZigZag + NativePType>(
 where
     <T as ExternalZigZag>::UInt: NativePType,
 {
-    PrimitiveArray::new(values.map_each(|v| T::encode(v)).freeze(), validity)
+    PrimitiveArray::new(
+        values.map_each_in_place(|v| T::encode(v)).freeze(),
+        validity,
+    )
 }
 
 pub fn zigzag_decode(parray: PrimitiveArray) -> PrimitiveArray {
@@ -57,7 +60,10 @@ fn zigzag_decode_primitive<T: ExternalZigZag + NativePType>(
 where
     <T as ExternalZigZag>::UInt: NativePType,
 {
-    PrimitiveArray::new(values.map_each(|v| T::decode(v)).freeze(), validity)
+    PrimitiveArray::new(
+        values.map_each_in_place(|v| T::decode(v)).freeze(),
+        validity,
+    )
 }
 
 #[cfg(test)]

encodings/zigzag/src/compress.rs

@@ -3,7 +3,8 @@
 
 use std::sync::Arc;
 
-use vortex_error::VortexResult;
+use vortex_dtype::DType;
+use vortex_error::{VortexResult, vortex_bail};
 use vortex_vector::Vector;
 
 use crate::execution::{BatchKernelRef, BindCtx};
@@ -33,6 +34,21 @@ pub trait ArrayOperator: 'static + Send + Sync {
 
 impl ArrayOperator for Arc<dyn Array> {
     fn execute_with_selection(&self, selection: Option<&ArrayRef>) -> VortexResult<Vector> {
+        if let Some(selection) = selection.as_ref() {
+            if !matches!(selection.dtype(), DType::Bool(_)) {
+                vortex_bail!(
+                    "Selection array must be of boolean type, got {}",
+                    selection.dtype()
+                );
+            }
+            if selection.len() != self.len() {
+                vortex_bail!(
+                    "Selection array length {} does not match array length {}",
+                    selection.len(),
+                    self.len()
+                );
+            }
+        }
         self.as_ref().execute_with_selection(selection)
     }
 

vortex-array/src/array/operator.rs

@@ -192,7 +192,7 @@ impl PrimitiveArray {
     {
         let validity = self.validity().clone();
         let buffer = match self.try_into_buffer_mut() {
-            Ok(buffer_mut) => buffer_mut.map_each(f),
+            Ok(buffer_mut) => buffer_mut.map_each_in_place(f),
             Err(parray) => BufferMut::<R>::from_iter(parray.buffer::<T>().iter().copied().map(f)),
         };
         PrimitiveArray::new(buffer.freeze(), validity)