Sparse pipeline driver

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

Author: gatesn

vortex-array/src/array/operator.rs

@@ -3,13 +3,13 @@
 
 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::driver::PipelineDriver;
 use crate::pipeline::PipelinedNode;
+use crate::pipeline::driver::PipelineDriver;
 use crate::vtable::{OperatorVTable, VTable};
 use crate::{Array, ArrayAdapter, ArrayRef};
 

vortex-array/src/pipeline/bit_view.rs

@@ -4,10 +4,11 @@
 use std::borrow::Cow;
 use std::fmt::{Debug, Formatter};
 
-use crate::pipeline::{N, N_BYTES, N_WORDS};
 use vortex_buffer::BitBuffer;
 use vortex_error::VortexResult;
 
+use crate::pipeline::{N, N_BYTES, N_WORDS};
+
 /// A borrowed fixed-size bit vector of length `N` bits, represented as an array of usize words.
 ///
 /// This struct is designed to provide a view over a Vortex [`vortex_buffer::BitBuffer`], therefore

vortex-array/src/pipeline/driver/allocation.rs

@@ -3,12 +3,13 @@
 
 //! Vector allocation strategy for pipelines
 
-use crate::pipeline::driver::{Node, NodeId};
-use crate::pipeline::{PipelineVector, VectorId, N};
-use crate::Array;
 use vortex_error::{VortexExpect, VortexResult};
 use vortex_vector::VectorMut;
 
+use crate::Array;
+use crate::pipeline::driver::{Node, NodeId};
+use crate::pipeline::{N, PipelineVector, VectorId};
+
 #[derive(Debug)]
 pub struct VectorAllocation {
     /// Where each node writes its output

vortex-array/src/pipeline/driver/bind.rs

@@ -1,12 +1,13 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use vortex_error::{VortexExpect, VortexResult};
+use vortex_vector::Vector;
+
 use crate::array::ArrayOperator;
-use crate::pipeline::driver::allocation::VectorAllocation;
 use crate::pipeline::driver::Node;
+use crate::pipeline::driver::allocation::VectorAllocation;
 use crate::pipeline::{BindContext, Kernel, VectorId};
-use vortex_error::{VortexExpect, VortexResult};
-use vortex_vector::Vector;
 
 pub(crate) fn bind_kernels(
     dag: &[Node],

vortex-array/src/pipeline/driver/mod.rs

@@ -3,25 +3,24 @@
 
 pub mod allocation;
 mod bind;
-mod input;
 mod toposort;
 
-use std::any::Any;
 use std::hash::{BuildHasher, Hash, Hasher};
 
-use crate::pipeline::bit_view::{BitView, BitViewExt};
-use crate::pipeline::driver::allocation::{allocate_vectors, OutputTarget};
-use crate::pipeline::driver::bind::bind_kernels;
-use crate::pipeline::driver::toposort::topological_sort;
-use crate::pipeline::{ElementPosition, Kernel, KernelCtx, PipelineInputs, PipelineVector, N};
-use crate::{Array, ArrayEq, ArrayHash, ArrayOperator, ArrayRef, ArrayVisitor, Precision};
 use itertools::Itertools;
-use vortex_dtype::{DType, NativePType};
-use vortex_error::{vortex_bail, VortexExpect, VortexResult};
+use vortex_dtype::DType;
+use vortex_error::{VortexResult, vortex_bail};
 use vortex_mask::Mask;
 use vortex_utils::aliases::hash_map::{HashMap, RandomState};
 use vortex_vector::{Vector, VectorMut, VectorMutOps};
 
+use crate::pipeline::bit_view::{BitView, BitViewExt};
+use crate::pipeline::driver::allocation::{OutputTarget, allocate_vectors};
+use crate::pipeline::driver::bind::bind_kernels;
+use crate::pipeline::driver::toposort::topological_sort;
+use crate::pipeline::{ElementPosition, Kernel, KernelCtx, N, PipelineInputs, PipelineVector};
+use crate::{Array, ArrayEq, ArrayHash, ArrayOperator, ArrayRef, ArrayVisitor, Precision};
+
 /// A pipeline driver takes a Vortex array and executes it into a canonical vector.
 ///
 /// The driver builds up a DAG of pipeline nodes from the array tree up to the edges of this
@@ -50,6 +49,7 @@ struct Node {
     // This node's underlying array.
     array: ArrayRef,
     /// The type of pipeline node.
+    #[allow(dead_code)] // TODO(ngates): pipeline execute does not yet use this
     kind: NodeKind,
     // The indices of the pipelined children nodes in the `nodes` vector.
     children: Vec<NodeId>,
@@ -68,8 +68,8 @@ enum NodeKind {
     /// A source node provides input to the pipeline by writing into mutable output vectors one
     /// batch at a time.
     Source,
+    /// A transform node takes pipelined inputs from its children and produces output vectors
     Transform,
-    Zip,
 }
 
 impl PipelineDriver {
@@ -247,21 +247,21 @@ impl Pipeline {
                 // The number of _full_ chunks we need to process.
                 let nchunks = selection.len() / N;
                 for _ in 0..nchunks {
-                    self.step(&mut self.ctx, &BitView::all_true(), &mut output)?;
+                    self.step(&BitView::all_true(), &mut output)?;
                 }
 
                 // Now process the final partial chunk, if any.
                 let remaining = selection.len() % N;
                 if remaining > 0 {
                     let selection_view = BitView::with_prefix(remaining);
-                    self.step(&mut self.ctx, &selection_view, &mut output)?;
+                    self.step(&selection_view, &mut output)?;
                 }
             }
             Mask::Values(mask_values) => {
                 // Loop over each chunk of N elements in the mask as a bit view.
                 let selection_bits = mask_values.bit_buffer();
                 for selection_view in selection_bits.iter_bit_views() {
-                    self.step(&mut self.ctx, &selection_view, &mut output)?;
+                    self.step(&selection_view, &mut output)?;
                 }
             }
         }
@@ -270,12 +270,7 @@ impl Pipeline {
     }
 
     /// Perform a single step of the pipeline.
-    fn step(
-        &self,
-        ctx: &mut KernelCtx,
-        selection: &BitView,
-        output: &mut VectorMut,
-    ) -> VortexResult<()> {
+    fn step(&mut self, selection: &BitView, output: &mut VectorMut) -> VortexResult<()> {
         // Loop over the kernels in toposorted execution order
         for &node_idx in self.exec_order.iter() {
             let kernel = &mut self.kernels[node_idx];
@@ -284,12 +279,12 @@ impl Pipeline {
             // take the intermediate vector and write into that.
             match &self.output_targets[node_idx] {
                 OutputTarget::ExternalOutput => {
-                    let output_len = output.len();
-                    let position = kernel.step(ctx, selection, output)?;
-                    if output.len() != output_len + N {
+                    let prev_output_len = output.len();
+                    let position = kernel.step(&self.ctx, selection, output)?;
+                    if output.len() != prev_output_len + N {
                         vortex_bail!(
                             "Kernel produced incorrect number of output elements, expected {}, got {}",
-                            output_len + N,
+                            prev_output_len + N,
                             output.len()
                         );
                     }
@@ -304,16 +299,15 @@ impl Pipeline {
                         ElementPosition::Compact => {
                             // The output is already compacted, we just need to adjust the length
                             // to cover only the selected elements.
-                            assert!(selection.true_count() <= N);
-                            unsafe { output.set_len(output_len + selection.true_count()) }
+                            output.truncate(prev_output_len + selection.true_count());
                         }
                     }
                 }
                 OutputTarget::IntermediateVector(vector_id) => {
-                    let mut out_vector = VectorMut::from(ctx.take_output(vector_id));
+                    let mut out_vector = VectorMut::from(self.ctx.take_output(vector_id));
                     out_vector.clear();
 
-                    let position = kernel.step(ctx, selection, &mut out_vector)?;
+                    let position = kernel.step(&self.ctx, selection, &mut out_vector)?;
                     if out_vector.len() != N {
                         vortex_bail!(
                             "Kernel produced incorrect number of output elements, expected {}, got {}",
@@ -325,7 +319,7 @@ impl Pipeline {
                     // Wrap the output vector back into a PipelineVector, indicating which position
                     // the elements are in.
                     let out_vector = PipelineVector::from_position(position, out_vector);
-                    ctx.replace_output(vector_id, out_vector)
+                    self.ctx.replace_output(vector_id, out_vector)
                 }
             };
         }

vortex-array/src/pipeline/driver/toposort.rs

@@ -3,7 +3,7 @@
 
 use std::collections::VecDeque;
 
-use vortex_error::{vortex_bail, VortexResult};
+use vortex_error::{VortexResult, vortex_bail};
 
 use crate::pipeline::driver::{Node, NodeId};
 

vortex-array/src/pipeline/mod.rs

@@ -4,13 +4,10 @@
 pub mod bit_view;
 pub mod driver;
 
-use std::ops::Deref;
-
 use vortex_error::{VortexExpect, VortexResult};
-use vortex_vector::{Vector, VectorMut, VectorMutOps};
+use vortex_vector::{Vector, VectorMut};
 
 use crate::pipeline::bit_view::BitView;
-use crate::Array;
 
 /// The number of elements in each step of a Vortex evaluation operator.
 pub const N: usize = 1024;
@@ -28,7 +25,7 @@ pub trait PipelinedNode {
     fn inputs(&self) -> PipelineInputs;
 
     /// Bind the node into a [`Kernel`] for pipelined execution.
-    fn bind(&self, ctx: &mut dyn BindContext) -> VortexResult<Box<dyn Kernel>>;
+    fn bind(&self, ctx: &dyn BindContext) -> VortexResult<Box<dyn Kernel>>;
 }
 
 /// Describes the type of pipeline node and its input information.
@@ -144,11 +141,6 @@ impl VectorId {
     fn new(idx: usize) -> Self {
         VectorId(idx)
     }
-
-    // Non-public getter to keep the type opaque to end users.
-    fn as_usize(&self) -> usize {
-        self.0
-    }
 }
 
 /// A pipeline vector passed into and out of pipeline kernels.

vortex-array/src/vtable/operator.rs

@@ -1,15 +1,15 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex_error::{vortex_bail, VortexResult};
+use vortex_error::{VortexResult, vortex_bail};
 use vortex_mask::Mask;
 use vortex_vector::Vector;
 
+use crate::ArrayRef;
 use crate::array::IntoArray;
 use crate::execution::{BatchKernelRef, BindCtx, ExecutionCtx};
 use crate::pipeline::PipelinedNode;
 use crate::vtable::{NotSupported, VTable};
-use crate::ArrayRef;
 
 /// A vtable for the new operator-based array functionality. Eventually this vtable will be
 /// merged into the main `VTable`, but for now it is kept separate to allow for incremental

vortex-buffer/src/bit/buf_mut.rs

@@ -267,6 +267,8 @@ impl BitBufferMut {
     }
 
     /// Truncate the buffer to the given length.
+    ///
+    /// If the given length is greater than the current length, this is a no-op.
     pub fn truncate(&mut self, len: usize) {
         if len > self.len {
             return;

vortex-vector/src/binaryview/vector_mut.rs

@@ -6,12 +6,12 @@
 use std::sync::Arc;
 
 use vortex_buffer::{BufferMut, ByteBuffer, ByteBufferMut};
-use vortex_error::{vortex_ensure, VortexExpect, VortexResult};
+use vortex_error::{VortexExpect, VortexResult, vortex_ensure};
 use vortex_mask::MaskMut;
 
-use crate::binaryview::vector::BinaryViewVector;
-use crate::binaryview::view::{validate_views, BinaryView};
 use crate::binaryview::BinaryViewType;
+use crate::binaryview::vector::BinaryViewVector;
+use crate::binaryview::view::{BinaryView, validate_views};
 use crate::{VectorMutOps, VectorOps};
 
 // Default capacity for new string data buffers of 2MiB.
@@ -207,6 +207,11 @@ impl<T: BinaryViewType> VectorMutOps for BinaryViewVectorMut<T> {
         self.open_buffer = None;
     }
 
+    fn truncate(&mut self, len: usize) {
+        self.views.truncate(len);
+        self.validity.truncate(len);
+    }
+
     fn extend_from_vector(&mut self, other: &BinaryViewVector<T>) {
         // Close any existing views into a new buffer
         self.flush_open_buffer();
@@ -261,7 +266,7 @@ mod tests {
     use std::ops::Deref;
     use std::sync::Arc;
 
-    use vortex_buffer::{buffer, buffer_mut, ByteBuffer};
+    use vortex_buffer::{ByteBuffer, buffer, buffer_mut};
     use vortex_mask::{Mask, MaskMut};
 
     use crate::binaryview::view::BinaryView;