Feature: batch execute for DeltaArray (#5667)

Tracking Issue: #5652

Implements `batch_execute` for `DeltaArray`.

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

Author: connortsui20

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

@@ -7,7 +7,6 @@ use fastlanes::Delta;
 use fastlanes::FastLanes;
 use fastlanes::Transpose;
 use num_traits::WrappingAdd;
-use vortex_array::Array;
 use vortex_array::ToCanonical;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::validity::Validity;
@@ -21,26 +20,32 @@ use crate::DeltaArray;
 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();
+
+    // TODO(connor): This is incorrect, we need to untranspose the validity!!!
+
+    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(crate) 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/kernel.rs

@@ -0,0 +1,103 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::fmt::Debug;
+use std::fmt::Formatter;
+
+use fastlanes::Delta;
+use fastlanes::FastLanes;
+use fastlanes::Transpose;
+use num_traits::WrappingAdd;
+use vortex_array::kernel::Kernel;
+use vortex_array::kernel::KernelRef;
+use vortex_array::kernel::PushDownResult;
+use vortex_dtype::NativePType;
+use vortex_dtype::PTypeDowncastExt;
+use vortex_error::VortexResult;
+use vortex_error::vortex_panic;
+use vortex_mask::Mask;
+use vortex_vector::Vector;
+use vortex_vector::primitive::PVector;
+use vortex_vector::primitive::PrimitiveVector;
+
+use crate::delta::array::delta_decompress::decompress_primitive;
+
+/// Holds the bound kernels and metadata needed to execute delta decompression.
+pub struct DeltaKernel {
+    pub(super) bases_kernel: KernelRef,
+    pub(super) deltas_kernel: KernelRef,
+    pub(super) start: usize,
+    pub(super) end: usize,
+    pub(super) validity: Mask,
+}
+
+impl Kernel for DeltaKernel {
+    fn execute(self: Box<Self>) -> VortexResult<Vector> {
+        // Extract all fields to avoid borrow issues.
+        let DeltaKernel {
+            bases_kernel,
+            deltas_kernel,
+            start,
+            end,
+            validity,
+        } = *self;
+
+        let bases = bases_kernel.execute()?.into_primitive();
+        let deltas = deltas_kernel.execute()?.into_primitive();
+
+        Ok(match bases {
+            PrimitiveVector::U8(pv) => {
+                decompress::<u8, { u8::LANES }>(&pv, &deltas, start, end, validity)
+            }
+            PrimitiveVector::U16(pv) => {
+                decompress::<u16, { u16::LANES }>(&pv, &deltas, start, end, validity)
+            }
+            PrimitiveVector::U32(pv) => {
+                decompress::<u32, { u32::LANES }>(&pv, &deltas, start, end, validity)
+            }
+            PrimitiveVector::U64(pv) => {
+                decompress::<u64, { u64::LANES }>(&pv, &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")
+            }
+        })
+    }
+
+    fn push_down_filter(self: Box<Self>, _selection: &Mask) -> VortexResult<PushDownResult> {
+        Ok(PushDownResult::NotPushed(self))
+    }
+}
+
+/// Decompresses delta-encoded data for a specific primitive type.
+fn decompress<T, const LANES: usize>(
+    bases: &PVector<T>,
+    deltas: &PrimitiveVector,
+    start: usize,
+    end: usize,
+    validity: Mask,
+) -> Vector
+where
+    T: NativePType + Delta + Transpose + WrappingAdd,
+{
+    let buffer = decompress_primitive::<T, LANES>(bases.as_ref(), deltas.downcast::<T>().as_ref());
+    let buffer = buffer.slice(start..end);
+
+    // SAFETY: We slice the buffer and the validity by the same range.
+    unsafe { PVector::<T>::new_unchecked(buffer, validity) }.into()
+}
+
+impl Debug for DeltaKernel {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("DeltaKernel")
+            .field("start", &self.start)
+            .field("end", &self.end)
+            .finish_non_exhaustive()
+    }
+}

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

@@ -5,6 +5,8 @@ use fastlanes::FastLanes;
 use prost::Message;
 use vortex_array::ArrayRef;
 use vortex_array::ProstMetadata;
+use vortex_array::kernel::BindCtx;
+use vortex_array::kernel::KernelRef;
 use vortex_array::serde::ArrayChildren;
 use vortex_array::vtable;
 use vortex_array::vtable::ArrayId;
@@ -21,10 +23,12 @@ use vortex_error::VortexResult;
 use vortex_error::vortex_ensure;
 use vortex_error::vortex_err;
 
+use self::kernel::DeltaKernel;
 use crate::DeltaArray;
 
 mod array;
 mod canonical;
+mod kernel;
 mod operations;
 mod validity;
 mod visitor;
@@ -117,6 +121,24 @@ impl VTable for DeltaVTable {
 
         DeltaArray::try_new(bases, deltas, metadata.0.offset as usize, len)
     }
+
+    fn bind_kernel(array: &DeltaArray, ctx: &mut BindCtx) -> VortexResult<KernelRef> {
+        let bases_kernel = array.bases().bind_kernel(ctx)?;
+        let deltas_kernel = array.deltas().bind_kernel(ctx)?;
+
+        let start = array.offset();
+        let end = start + array.len();
+
+        let validity = array.deltas().validity_mask().slice(start..end);
+
+        Ok(Box::new(DeltaKernel {
+            bases_kernel,
+            deltas_kernel,
+            start,
+            end,
+            validity,
+        }))
+    }
 }
 
 #[derive(Debug)]