batch execute for delta array

Signed-off-by: Connor Tsui <[email protected]>

Author: connortsui20

encodings/fastlanes/src/delta/array/delta_decompress.rs

@@ -10,37 +10,92 @@ use num_traits::WrappingAdd;
 use vortex_array::Array;
 use vortex_array::ToCanonical;
 use vortex_array::arrays::PrimitiveArray;
+use vortex_array::execution::ExecutionCtx;
 use vortex_array::validity::Validity;
 use vortex_buffer::Buffer;
 use vortex_buffer::BufferMut;
 use vortex_dtype::NativePType;
+use vortex_dtype::PTypeDowncastExt;
 use vortex_dtype::match_each_unsigned_integer_ptype;
+use vortex_error::VortexResult;
+use vortex_error::vortex_panic;
+use vortex_vector::Vector;
+use vortex_vector::primitive::PVector;
+use vortex_vector::primitive::PrimitiveVector;
 
 use crate::DeltaArray;
 
+/// Decompresses the [`DeltaArray`] into a vector.
+pub fn delta_decompress_into_primitive_vector(
+    array: &DeltaArray,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<Vector> {
+    let bases = array.bases().batch_execute(ctx)?.into_primitive();
+    let deltas = array.deltas().batch_execute(ctx)?.into_primitive();
+
+    let start = array.offset();
+    let end = start + array.len();
+
+    let validity = array.deltas().validity_mask().slice(start..end);
+
+    macro_rules! decompress_case {
+        ($pv:expr, $type:ty, $deltas:expr, $start:expr, $end:expr, $validity:expr) => {{
+            const LANES: usize = <$type>::LANES;
+            let buffer = decompress_primitive::<$type, LANES>(
+                $pv.as_ref(),
+                $deltas.downcast::<$type>().as_ref(),
+            );
+            let buffer = buffer.slice($start..$end);
+
+            // SAFETY: We slice the buffer and the validity by the same range.
+            unsafe { PVector::<$type>::new_unchecked(buffer, $validity) }.into()
+        }};
+    }
+
+    Ok(match bases {
+        PrimitiveVector::U8(pv) => decompress_case!(pv, u8, deltas, start, end, validity),
+        PrimitiveVector::U16(pv) => decompress_case!(pv, u16, deltas, start, end, validity),
+        PrimitiveVector::U32(pv) => decompress_case!(pv, u32, deltas, start, end, validity),
+        PrimitiveVector::U64(pv) => decompress_case!(pv, u64, deltas, start, end, validity),
+        PrimitiveVector::I8(_)
+        | PrimitiveVector::I16(_)
+        | PrimitiveVector::I32(_)
+        | PrimitiveVector::I64(_)
+        | PrimitiveVector::F16(_)
+        | PrimitiveVector::F32(_)
+        | PrimitiveVector::F64(_) => {
+            vortex_panic!("Tried to match a non-unsigned vector in an unsigned match statement")
+        }
+    })
+}
+
 pub fn delta_decompress(array: &DeltaArray) -> PrimitiveArray {
     let bases = array.bases().to_primitive();
     let deltas = array.deltas().to_primitive();
-    let decoded = match_each_unsigned_integer_ptype!(deltas.ptype(), |T| {
+
+    let start = array.offset();
+    let end = start + array.len();
+
+    let validity = Validity::from_mask(array.deltas().validity_mask(), array.dtype().nullability());
+    let validity = validity.slice(start..end);
+
+    match_each_unsigned_integer_ptype!(deltas.ptype(), |T| {
         const LANES: usize = T::LANES;
 
-        PrimitiveArray::new(
-            decompress_primitive::<T, LANES>(bases.as_slice(), deltas.as_slice()),
-            Validity::from_mask(array.deltas().validity_mask(), array.dtype().nullability()),
-        )
-    });
+        let buffer = decompress_primitive::<T, LANES>(bases.as_slice(), deltas.as_slice());
+        let buffer = buffer.slice(start..end);
 
-    decoded
-        .slice(array.offset()..array.offset() + array.len())
-        .to_primitive()
+        PrimitiveArray::new(buffer, validity)
+    })
 }
 
 // TODO(ngates): can we re-use the deltas buffer for the result? Might be tricky given the
 //  traversal ordering, but possibly doable.
-fn decompress_primitive<T: NativePType + Delta + Transpose + WrappingAdd, const LANES: usize>(
-    bases: &[T],
-    deltas: &[T],
-) -> Buffer<T> {
+/// Performs the low-level delta decompression on primitive values.
+pub(super) fn decompress_primitive<T, const LANES: usize>(bases: &[T], deltas: &[T]) -> Buffer<T>
+where
+    T: NativePType + Delta + Transpose + WrappingAdd,
+{
     // How many fastlanes vectors we will process.
     let num_chunks = deltas.len() / 1024;
 

encodings/fastlanes/src/delta/vtable/mod.rs

@@ -4,6 +4,7 @@
 use fastlanes::FastLanes;
 use prost::Message;
 use vortex_array::ProstMetadata;
+use vortex_array::execution::ExecutionCtx;
 use vortex_array::serde::ArrayChildren;
 use vortex_array::vtable;
 use vortex_array::vtable::ArrayId;
@@ -18,8 +19,10 @@ use vortex_dtype::PType;
 use vortex_dtype::match_each_unsigned_integer_ptype;
 use vortex_error::VortexResult;
 use vortex_error::vortex_err;
+use vortex_vector::Vector;
 
 use crate::DeltaArray;
+use crate::delta::array::delta_decompress::delta_decompress_into_primitive_vector;
 
 mod array;
 mod canonical;
@@ -98,6 +101,10 @@ impl VTable for DeltaVTable {
 
         DeltaArray::try_new(bases, deltas, metadata.0.offset as usize, len)
     }
+
+    fn batch_execute(array: &DeltaArray, ctx: &mut ExecutionCtx) -> VortexResult<Vector> {
+        delta_decompress_into_primitive_vector(array, ctx)
+    }
 }
 
 #[derive(Debug)]