batch execute for delta array

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

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

@@ -15,28 +15,55 @@ use vortex_buffer::Buffer;
 use vortex_buffer::BufferMut;
 use vortex_dtype::NativePType;
 use vortex_dtype::match_each_unsigned_integer_ptype;
+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) -> PrimitiveVector {
+    let bases = array.bases().to_primitive();
+    let deltas = array.deltas().to_primitive();
+
+    let start = array.offset();
+    let end = start + array.len();
+
+    let validity = array.deltas().validity_mask().slice(start..end);
+
+    match_each_unsigned_integer_ptype!(deltas.ptype(), |T| {
+        const LANES: usize = T::LANES;
+
+        let buffer = decompress_primitive::<T, LANES>(bases.as_slice(), deltas.as_slice());
+        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()
+    })
+}
+
 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.
+/// Performs the low-level delta decompression on primitive values.
 fn decompress_primitive<T: NativePType + Delta + Transpose + WrappingAdd, const LANES: usize>(
     bases: &[T],
     deltas: &[T],
@@ -84,3 +111,46 @@ fn decompress_primitive<T: NativePType + Delta + Transpose + WrappingAdd, const
 
     output.freeze()
 }
+
+#[cfg(test)]
+mod tests {
+    use vortex_array::arrays::PrimitiveArray;
+    use vortex_dtype::PTypeDowncast;
+    use vortex_vector::VectorOps;
+
+    use super::*;
+    use crate::DeltaArray;
+
+    #[test]
+    fn test_decompress_into_vector() {
+        let input: PrimitiveArray = (0u32..10_000).collect();
+        let delta = DeltaArray::try_from_primitive_array(&input).unwrap();
+        assert_eq!(delta.len(), input.len());
+
+        let decompressed = delta_decompress_into_primitive_vector(&delta);
+        let decompressed = decompressed.into_u32();
+
+        assert_eq!(decompressed.len(), input.len());
+        for (actual, expected) in decompressed.as_ref().iter().zip(input.as_slice::<u32>()) {
+            assert_eq!(actual, expected);
+        }
+        assert_eq!(decompressed.validity(), &input.validity_mask());
+    }
+
+    #[test]
+    fn test_decompress_into_vector_nullable() {
+        let input =
+            PrimitiveArray::from_option_iter((0u32..10_000).map(|i| (i % 2 == 0).then_some(i)));
+        let delta = DeltaArray::try_from_primitive_array(&input).unwrap();
+        assert_eq!(delta.len(), input.len());
+
+        let decompressed = delta_decompress_into_primitive_vector(&delta);
+        let decompressed = decompressed.into_u32();
+
+        assert_eq!(decompressed.len(), input.len());
+        for (actual, expected) in decompressed.as_ref().iter().zip(input.as_slice::<u32>()) {
+            assert_eq!(actual, expected);
+        }
+        assert_eq!(decompressed.validity(), &input.validity_mask());
+    }
+}

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> {
+        Ok(delta_decompress_into_primitive_vector(array).into())
+    }
 }
 
 #[derive(Debug)]