pipelined execution

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

Author: gatesn

vortex-array/src/array/operator.rs

@@ -3,9 +3,9 @@
 
 use std::sync::Arc;
 
-use vortex_error::{vortex_panic, VortexResult};
+use vortex_error::{VortexResult, vortex_panic};
 use vortex_mask::Mask;
-use vortex_vector::{vector_matches_dtype, Vector, VectorOps};
+use vortex_vector::{Vector, VectorOps, vector_matches_dtype};
 
 use crate::execution::{BatchKernelRef, BindCtx, DummyExecutionCtx, ExecutionCtx};
 use crate::pipeline::source_driver::PipelineSourceDriver;

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

@@ -1,18 +1,20 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex_buffer::Buffer;
+use vortex_buffer::{BitBuffer, Buffer};
 use vortex_compute::filter::Filter;
-use vortex_dtype::{match_each_native_ptype, NativePType, PTypeDowncastExt};
+use vortex_dtype::{NativePType, PTypeDowncastExt, match_each_native_ptype};
 use vortex_error::VortexResult;
-use vortex_mask::Mask;
 use vortex_vector::primitive::PVector;
 use vortex_vector::{VectorMut, VectorMutOps};
 
 use crate::arrays::{MaskedVTable, PrimitiveArray, PrimitiveVTable};
-use crate::execution::{kernel, BatchKernelRef, BindCtx};
-use crate::pipeline::bit_view::BitView;
-use crate::pipeline::{BindContext, KernelContext, PipelinedSource, SourceKernel, N};
+use crate::execution::{BatchKernelRef, BindCtx, kernel};
+use crate::pipeline::bit_view::{BitSlice, BitView};
+use crate::pipeline::{
+    AllNullSourceKernel, BindContext, KernelContext, N, PipelinedSource, SourceKernel,
+};
+use crate::validity::Validity;
 use crate::vtable::{OperatorVTable, ValidityHelper};
 use crate::{ArrayRef, IntoArray};
 
@@ -71,25 +73,81 @@ impl OperatorVTable<PrimitiveVTable> for PrimitiveVTable {
 }
 
 impl PipelinedSource for PrimitiveArray {
-    fn bind_source(&self, _ctx: &mut dyn BindContext) -> VortexResult<Box<dyn SourceKernel>> {
-        match_each_native_ptype!(self.ptype(), |T| {
-            let primitive_kernel = PrimitiveKernel {
-                buffer: self.buffer::<T>().clone(),
-                validity: self.validity_mask(),
-                offset: 0,
-            };
-            Ok(Box::new(primitive_kernel))
-        })
+    fn bind_source(&self, ctx: &mut dyn BindContext) -> VortexResult<Box<dyn SourceKernel>> {
+        match self.validity() {
+            Validity::NonNullable | Validity::AllValid => {
+                match_each_native_ptype!(self.ptype(), |T| {
+                    let primitive_kernel = NonNullablePrimitiveKernel {
+                        buffer: self.buffer::<T>(),
+                        offset: 0,
+                    };
+                    Ok(Box::new(primitive_kernel))
+                })
+            }
+            Validity::AllInvalid => Ok(Box::new(AllNullSourceKernel)),
+            Validity::Array(_) => {
+                let validity = ctx.batch_input(0).into_bool();
+                // Validity is non-nullable, so we extract the inner bit buffer.
+                let (validity, _) = validity.into_parts();
+
+                match_each_native_ptype!(self.ptype(), |T| {
+                    let primitive_kernel = NullablePrimitiveKernel {
+                        buffer: self.buffer::<T>(),
+                        validity,
+                        offset: 0,
+                    };
+                    Ok(Box::new(primitive_kernel))
+                })
+            }
+        }
+    }
+}
+
+struct NonNullablePrimitiveKernel<T: NativePType> {
+    buffer: Buffer<T>,
+    offset: usize,
+}
+
+impl<T: NativePType> SourceKernel for NonNullablePrimitiveKernel<T> {
+    fn skip(&mut self, n: usize) {
+        self.offset += n * N;
+    }
+
+    fn step(
+        &mut self,
+        _ctx: &KernelContext,
+        selection: &BitView,
+        out: &mut VectorMut,
+    ) -> VortexResult<()> {
+        let out = out.as_primitive_mut().downcast::<T>();
+
+        // SAFETY: we know the output has sufficient capacity.
+        unsafe {
+            out.validity_mut().append_n(true, 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(|BitSlice { start, len }| {
+            out.as_mut()[out_pos..][..len].copy_from_slice(&source[start..][..len]);
+            out_pos += len;
+        });
+
+        Ok(())
     }
 }
 
-struct PrimitiveKernel<T: NativePType> {
+struct NullablePrimitiveKernel<T: NativePType> {
     buffer: Buffer<T>,
-    validity: Mask,
+    #[allow(dead_code)] // TODO(ngates): implement appending validity bits
+    validity: BitBuffer,
     offset: usize,
 }
 
-impl<T: NativePType> SourceKernel for PrimitiveKernel<T> {
+impl<T: NativePType> SourceKernel for NullablePrimitiveKernel<T> {
     fn skip(&mut self, n: usize) {
         self.offset += n * N;
     }
@@ -114,9 +172,14 @@ impl<T: NativePType> SourceKernel for PrimitiveKernel<T> {
         let source = &self.buffer.as_slice()[self.offset..];
 
         let mut out_pos = 0;
-        selection.iter_slices(|(start, len)| {
+        selection.iter_slices(|BitSlice { start, len }| {
+            // Copy over the elements.
             out.as_mut()[out_pos..][..len].copy_from_slice(&source[start..][..len]);
             out_pos += len;
+
+            // Append the validity bits.
+            let _validity = unsafe { out.validity_mut() };
+            todo!("Append validity bits correctly and optimally!");
         });
 
         Ok(())

vortex-array/src/pipeline/bit_view.rs

@@ -212,11 +212,13 @@ impl<'a> BitView<'a> {
         }
     }
 
+    /// Runs the provided function `f` for each range of `true` bits in the view.
+    ///
+    /// The function `f` receives a [`BitSlice`] containing the inclusive `start` bit as well as
+    /// the length.
     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)),
+        F: FnMut(BitSlice),
     {
         if self.true_count == 0 {
             return;
@@ -231,7 +233,10 @@ impl<'a> BitView<'a> {
                 0 => {
                     // If a slice was being tracked, the run ends at the start of this word.
                     if slice_length > 0 {
-                        f((slice_start_bit, slice_length));
+                        f(BitSlice {
+                            start: slice_start_bit,
+                            len: slice_length,
+                        });
                         slice_length = 0;
                     }
                 }
@@ -250,7 +255,10 @@ impl<'a> BitView<'a> {
 
                         // If a run was open, and we hit a zero gap, report the finished slice
                         if slice_length > 0 && zeros > 0 {
-                            f((slice_start_bit, slice_length));
+                            f(BitSlice {
+                                start: slice_start_bit,
+                                len: slice_length,
+                            });
                             slice_length = 0; // Reset state for a new slice
                         }
 
@@ -284,22 +292,26 @@ impl<'a> BitView<'a> {
         }
 
         if slice_length > 0 {
-            f((slice_start_bit, slice_length));
+            f(BitSlice {
+                start: slice_start_bit,
+                len: slice_length,
+            });
         }
     }
 
-    /// Runs the provided function `f` for each range of `true` bits in the view.
-    ///
-    /// The function `f` receives a tuple `(start, len)` where `start` is the index of the first
-    /// `true` bit and `len` is the number of consecutive `true` bits.
-    ///
-    /// FIXME(ngates): this code is broken.
-
     pub fn as_raw(&self) -> &[u8; N_BYTES] {
         self.bits.as_ref()
     }
 }
 
+/// A slice of bits within a [`BitBuffer`].
+///
+/// We use this struct to avoid a common mistake of assuming the slices represent (start, end) ranges,
+pub struct BitSlice {
+    pub start: usize,
+    pub len: usize,
+}
+
 pub trait BitViewExt {
     /// Iterate the [`BitBuffer`] in fixed-size chunks of [`BitView`].
     ///
@@ -319,7 +331,7 @@ impl BitViewExt for BitBuffer {
             0,
             "BitView iteration requires zero bit offset"
         );
-        let n_views = (self.len() + N - 1) / N;
+        let n_views = self.len().div_ceil(N);
         BitViewIterator {
             bits: self.inner().as_ref(),
             view_idx: 0,
@@ -365,8 +377,6 @@ impl<'a> Iterator for BitViewIterator<'a> {
 
 #[cfg(test)]
 mod tests {
-    use std::usize;
-
     use super::*;
 
     #[test]
@@ -591,7 +601,6 @@ mod tests {
             view.iter_slices(|slice| slices.push(slice));
 
             assert_eq!(slices.len(), 1);
-            assert_eq!(slices[0], (0, i));
         }
     }
 }

vortex-array/src/pipeline/mod.rs

@@ -25,7 +25,7 @@ use std::ops::Deref;
 
 use bit_view::BitView;
 use vortex_error::VortexResult;
-use vortex_vector::{Vector, VectorMut};
+use vortex_vector::{Vector, VectorMut, VectorMutOps};
 
 use crate::Array;
 
@@ -150,3 +150,20 @@ impl KernelContext {
         &self.vectors[vector_id]
     }
 }
+
+/// A general implementation of a source kernel that produces all null values.
+pub struct AllNullSourceKernel;
+
+impl SourceKernel for AllNullSourceKernel {
+    fn skip(&mut self, _n: usize) {}
+
+    fn step(
+        &mut self,
+        _ctx: &KernelContext,
+        selection: &BitView,
+        out: &mut VectorMut,
+    ) -> VortexResult<()> {
+        out.append_nulls(selection.true_count());
+        Ok(())
+    }
+}

vortex-array/src/pipeline/source_driver.rs

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

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::{vortex_panic, VortexExpect};
+use vortex_error::{VortexExpect, vortex_panic};
 
 use crate::debug::TruncatedDebug;
 use crate::trusted_len::TrustedLen;
@@ -731,7 +731,7 @@ impl Write for ByteBufferMut {
 mod test {
     use bytes::{Buf, BufMut};
 
-    use crate::{buffer_mut, Alignment, BufferMut, ByteBufferMut};
+    use crate::{Alignment, BufferMut, ByteBufferMut, buffer_mut};
 
     #[test]
     fn capacity() {

vortex-mask/src/mask_mut.rs

@@ -95,6 +95,10 @@ impl MaskMut {
     }
 
     /// Set the length of the mask.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that `new_len` is less than the capacity of the mask.
     pub unsafe fn set_len(&mut self, new_len: usize) {
         debug_assert!(new_len < self.capacity());
         match &mut self.0 {