pipelined execution

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

Author: gatesn

vortex-array/src/array/operator.rs

@@ -3,11 +3,12 @@
 
 use std::sync::Arc;
 
-use vortex_error::{VortexResult, vortex_panic};
+use vortex_error::{vortex_panic, VortexResult};
 use vortex_mask::Mask;
-use vortex_vector::{Vector, VectorOps, vector_matches_dtype};
+use vortex_vector::{vector_matches_dtype, Vector, VectorOps};
 
 use crate::execution::{BatchKernelRef, BindCtx, DummyExecutionCtx, ExecutionCtx};
+use crate::pipeline::source_driver::PipelineSourceDriver;
 use crate::vtable::{OperatorVTable, VTable};
 use crate::{Array, ArrayAdapter, ArrayRef};
 
@@ -62,6 +63,13 @@ impl ArrayOperator for Arc<dyn Array> {
 
 impl<V: VTable> ArrayOperator for ArrayAdapter<V> {
     fn execute_batch(&self, selection: &Mask, ctx: &mut dyn ExecutionCtx) -> VortexResult<Vector> {
+        // Check if the array is a pipeline source, and if so use the single-node driver for now.
+        if let Some(pipeline_source) =
+            <V::OperatorVTable as OperatorVTable<V>>::as_pipelined_source(&self.0)
+        {
+            return PipelineSourceDriver::new(pipeline_source).execute(selection);
+        }
+
         let vector =
             <V::OperatorVTable as OperatorVTable<V>>::execute_batch(&self.0, selection, ctx)?;
 

vortex-array/src/arrays/primitive/vtable/operator.rs

@@ -3,16 +3,16 @@
 
 use vortex_buffer::Buffer;
 use vortex_compute::filter::Filter;
-use vortex_dtype::{NativePType, PTypeDowncastExt, match_each_native_ptype};
+use vortex_dtype::{match_each_native_ptype, NativePType, PTypeDowncastExt};
 use vortex_error::VortexResult;
 use vortex_mask::Mask;
-use vortex_vector::VectorMut;
 use vortex_vector::primitive::PVector;
+use vortex_vector::{VectorMut, VectorMutOps};
 
 use crate::arrays::{MaskedVTable, PrimitiveArray, PrimitiveVTable};
-use crate::execution::{BatchKernelRef, BindCtx, kernel};
+use crate::execution::{kernel, BatchKernelRef, BindCtx};
 use crate::pipeline::bit_view::BitView;
-use crate::pipeline::{BindContext, KernelContext, N, PipelinedSource, SourceKernel};
+use crate::pipeline::{BindContext, KernelContext, PipelinedSource, SourceKernel, N};
 use crate::vtable::{OperatorVTable, ValidityHelper};
 use crate::{ArrayRef, IntoArray};
 
@@ -106,16 +106,16 @@ impl<T: NativePType> SourceKernel for PrimitiveKernel<T> {
         //  separately from copying over the elements.
         unsafe {
             out.validity_mut().append_n(true, selection.true_count());
-            out.elements_mut().set_len(selection.true_count());
+            let prev_len = out.len();
+            out.elements_mut()
+                .set_len(prev_len + selection.true_count());
         }
 
         let source = &self.buffer.as_slice()[self.offset..];
 
         let mut out_pos = 0;
-        selection.iter_slices(|(start, end)| {
-            print!("Slicing {} to {}\n", start, end);
-            let len = end - start;
-            out.as_mut()[out_pos..][..len].copy_from_slice(&source[start..end]);
+        selection.iter_slices(|(start, len)| {
+            out.as_mut()[out_pos..][..len].copy_from_slice(&source[start..][..len]);
             out_pos += len;
         });
 

vortex-array/src/pipeline/bit_view.rs

@@ -214,6 +214,8 @@ impl<'a> BitView<'a> {
 
     pub fn iter_slices<F>(&self, mut f: F)
     where
+        // FIXME(ngates): I have repeatedly assumed this to be a (start, end) slice, not a
+        //  (start, len)... I think we should wrap this in a struct to avoid confusion.
         F: FnMut((usize, usize)),
     {
         if self.true_count == 0 {

vortex-array/src/pipeline/source_driver.rs

@@ -2,14 +2,14 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use itertools::Itertools;
-use vortex_error::{VortexResult, vortex_panic};
+use vortex_error::{vortex_panic, VortexResult};
 use vortex_mask::Mask;
 use vortex_vector::{Vector, VectorMut, VectorMutOps};
 
 use crate::pipeline::bit_view::{BitView, BitViewExt};
-use crate::pipeline::{BindContext, KernelContext, N, PipelinedSource, VectorId};
+use crate::pipeline::{BindContext, KernelContext, PipelinedSource, VectorId, N};
 
-/// Temporary driver for executing a single array in a pipelined fashion.
+/// Temporary driver for executing a single source array in a pipelined fashion.
 pub struct PipelineSourceDriver<'a> {
     array: &'a dyn PipelinedSource,
 }

vortex-buffer/src/buffer_mut.rs

@@ -9,7 +9,7 @@ use std::ops::{Deref, DerefMut};
 
 use bytes::buf::UninitSlice;
 use bytes::{Buf, BufMut, BytesMut};
-use vortex_error::{VortexExpect, vortex_panic};
+use vortex_error::{vortex_panic, VortexExpect};
 
 use crate::debug::TruncatedDebug;
 use crate::trusted_len::TrustedLen;
@@ -241,10 +241,15 @@ impl<T> BufferMut<T> {
         }
     }
 
+    /// Sets the length of the buffer.
+    ///
     /// # Safety
-    /// The caller must ensure that the buffer was properly initialized up to `len`.
+    ///
+    /// The caller must ensure that there is sufficient capacity in the buffer and that the values
+    /// are valid up to `len`.
     #[inline]
     pub unsafe fn set_len(&mut self, len: usize) {
+        debug_assert!(len <= self.capacity());
         unsafe { self.bytes.set_len(len * size_of::<T>()) };
         self.length = len;
     }
@@ -726,7 +731,7 @@ impl Write for ByteBufferMut {
 mod test {
     use bytes::{Buf, BufMut};
 
-    use crate::{Alignment, BufferMut, ByteBufferMut, buffer_mut};
+    use crate::{buffer_mut, Alignment, BufferMut, ByteBufferMut};
 
     #[test]
     fn capacity() {