API and partial implementation of sparse pipelined exection (#5284)

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

Author: gatesn

vortex-array/src/array/operator.rs

@@ -8,8 +8,9 @@ use vortex_mask::Mask;
 use vortex_vector::{Vector, VectorOps, vector_matches_dtype};
 
 use crate::execution::{BatchKernelRef, BindCtx, DummyExecutionCtx, ExecutionCtx};
-use crate::pipeline::source_driver::PipelineSourceDriver;
-use crate::vtable::{OperatorVTable, PipelineNode, VTable};
+use crate::pipeline::PipelinedNode;
+use crate::pipeline::driver::PipelineDriver;
+use crate::vtable::{OperatorVTable, VTable};
 use crate::{Array, ArrayAdapter, ArrayRef};
 
 /// Array functions as provided by the `OperatorVTable`.
@@ -31,6 +32,9 @@ pub trait ArrayOperator: 'static + Send + Sync {
     /// Optimize the array by pushing down a parent array.
     fn reduce_parent(&self, parent: &ArrayRef, child_idx: usize) -> VortexResult<Option<ArrayRef>>;
 
+    /// Returns the array as a pipeline node, if supported.
+    fn as_pipelined(&self) -> Option<&dyn PipelinedNode>;
+
     /// Bind the array to a batch kernel. This is an internal function
     fn bind(
         &self,
@@ -52,6 +56,10 @@ impl ArrayOperator for Arc<dyn Array> {
         self.as_ref().reduce_parent(parent, child_idx)
     }
 
+    fn as_pipelined(&self) -> Option<&dyn PipelinedNode> {
+        self.as_ref().as_pipelined()
+    }
+
     fn bind(
         &self,
         selection: Option<&ArrayRef>,
@@ -63,14 +71,6 @@ 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 node
-        if let Some(pipeline_node) =
-            <V::OperatorVTable as OperatorVTable<V>>::pipeline_node(&self.0)
-            && let PipelineNode::Source(source) = pipeline_node
-        {
-            return PipelineSourceDriver::new(source).execute(selection);
-        }
-
         let vector =
             <V::OperatorVTable as OperatorVTable<V>>::execute_batch(&self.0, selection, ctx)?;
 
@@ -104,6 +104,10 @@ impl<V: VTable> ArrayOperator for ArrayAdapter<V> {
         <V::OperatorVTable as OperatorVTable<V>>::reduce_parent(&self.0, parent, child_idx)
     }
 
+    fn as_pipelined(&self) -> Option<&dyn PipelinedNode> {
+        <V::OperatorVTable as OperatorVTable<V>>::pipeline_node(&self.0)
+    }
+
     fn bind(
         &self,
         selection: Option<&ArrayRef>,
@@ -126,11 +130,17 @@ impl BindCtx for () {
 
 impl dyn Array + '_ {
     pub fn execute(&self) -> VortexResult<Vector> {
-        self.execute_batch(&Mask::new_true(self.len()), &mut DummyExecutionCtx)
+        self.execute_with_selection(&Mask::new_true(self.len()))
     }
 
-    pub fn execute_with_selection(&self, mask: &Mask) -> VortexResult<Vector> {
-        assert_eq!(self.len(), mask.len());
-        self.execute_batch(mask, &mut DummyExecutionCtx)
+    pub fn execute_with_selection(&self, selection: &Mask) -> VortexResult<Vector> {
+        assert_eq!(self.len(), selection.len());
+
+        // Check if the array is a pipeline node
+        if self.as_pipelined().is_some() {
+            return PipelineDriver::new(self.to_array()).execute(selection);
+        }
+
+        self.execute_batch(selection, &mut DummyExecutionCtx)
     }
 }

…x-array/src/pipeline/operator/buffers.rs …-array/src/pipeline/driver/allocation.rsvortex-array/src/pipeline/operator/buffers.rs renamed to vortex-array/src/pipeline/driver/allocation.rs vortex-array/src/pipeline/operator/buffers.rs renamed to vortex-array/src/pipeline/driver/allocation.rs

@@ -3,22 +3,19 @@
 
 //! Vector allocation strategy for pipelines
 
-#![allow(dead_code)]
-
-use std::cell::RefCell;
-
 use vortex_error::{VortexExpect, VortexResult};
+use vortex_vector::VectorMut;
 
-use crate::pipeline::operator::{NodeId, PipelineNode};
-use crate::pipeline::vec::Vector;
-use crate::pipeline::{VType, VectorId};
+use crate::Array;
+use crate::pipeline::driver::{Node, NodeId};
+use crate::pipeline::{N, PipelineVector, VectorId};
 
 #[derive(Debug)]
-pub struct VectorAllocationPlan {
+pub struct VectorAllocation {
     /// Where each node writes its output
     pub(crate) output_targets: Vec<OutputTarget>,
     /// The actual allocated vectors
-    pub(crate) vectors: Vec<RefCell<Vector>>,
+    pub(crate) vectors: Vec<PipelineVector>,
 }
 
 // TODO(joe): support in-place view operations
@@ -30,41 +27,34 @@ pub(crate) enum OutputTarget {
     /// Node writes to the top-level provided output
     ExternalOutput,
     /// Node writes to an allocated intermediate vector
-    IntermediateVector(usize), // vector idx
+    IntermediateVector(VectorId), // vector idx
 }
 
 impl OutputTarget {
     pub fn vector_id(&self) -> Option<VectorId> {
         match self {
-            OutputTarget::IntermediateVector(idx) => Some(*idx),
+            OutputTarget::IntermediateVector(vector_id) => Some(*vector_id),
             OutputTarget::ExternalOutput => None,
         }
     }
 }
 
-/// Represents an allocated vector that can be reused
-#[derive(Debug, Clone)]
-struct VectorAllocation {
-    /// Type of elements in this vector
-    element_type: VType,
-}
-
 // ============================================================================
 // Improved Pipeline with vector allocation
 // ============================================================================
 
 /// Allocate vectors with lifetime analysis and zero-copy optimization
 pub(super) fn allocate_vectors(
-    dag: &[PipelineNode],
+    dag: &[Node],
     execution_order: &[NodeId],
-) -> VortexResult<VectorAllocationPlan> {
+) -> VortexResult<VectorAllocation> {
     let mut output_targets: Vec<Option<OutputTarget>> = vec![None; dag.len()];
-    let mut allocations = Vec::new();
+    let mut allocation_types = Vec::new();
 
     // Process nodes in reverse execution order (top-down for output propagation)
     for &node_idx in execution_order.iter().rev() {
         let node = &dag[node_idx];
-        let operator = &node.operator;
+        let array = &node.array;
 
         // Determine output target
         let output_target = if node.parents.is_empty() {
@@ -81,24 +71,22 @@ pub(super) fn allocate_vectors(
             // 3. Verify no conflicts with parallel execution paths
             // 4. Ensure proper vector lifetime management
 
-            let alloc_id = allocations.len();
-            allocations.push(VectorAllocation {
-                element_type: operator.dtype().into(),
-            });
-            OutputTarget::IntermediateVector(alloc_id)
+            let vector_id = VectorId::new(allocation_types.len());
+            allocation_types.push(array.dtype());
+            OutputTarget::IntermediateVector(vector_id)
         };
 
         output_targets[node_idx] = Some(output_target);
     }
 
-    Ok(VectorAllocationPlan {
+    Ok(VectorAllocation {
         output_targets: output_targets
             .into_iter()
             .map(|target| target.vortex_expect("missing target"))
             .collect(),
-        vectors: allocations
+        vectors: allocation_types
             .into_iter()
-            .map(|alloc| RefCell::new(Vector::new_with_vtype(alloc.element_type)))
+            .map(|dtype| PipelineVector::Compact(VectorMut::with_capacity(dtype, 2 * N)))
             .collect(),
     })
 }

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

@@ -0,0 +1,63 @@
+// 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::Node;
+use crate::pipeline::driver::allocation::VectorAllocation;
+use crate::pipeline::{BindContext, Kernel, VectorId};
+
+pub(crate) fn bind_kernels(
+    dag: &[Node],
+    allocation_plan: &VectorAllocation,
+    all_batch_inputs: &[Vector],
+) -> VortexResult<Vec<Box<dyn Kernel>>> {
+    let mut kernels = Vec::with_capacity(dag.len());
+    for node in dag {
+        let input_ids = node
+            .children
+            .iter()
+            .map(|node_id| {
+                allocation_plan.output_targets[*node_id]
+                    .vector_id()
+                    .vortex_expect("Input node must have an output vector ID")
+            })
+            .collect::<Vec<_>>();
+
+        let batch_inputs: Vec<_> = node
+            .batch_inputs
+            .iter()
+            .map(|idx| all_batch_inputs[*idx].clone())
+            .collect();
+
+        let bind_context = PipelineBindContext {
+            children: &input_ids,
+            batch_inputs: &batch_inputs,
+        };
+
+        let pipelined = node
+            .array
+            .as_pipelined()
+            .vortex_expect("Array in pipeline DAG does not support pipelined execution");
+
+        kernels.push(pipelined.bind(&bind_context)?);
+    }
+    Ok(kernels)
+}
+
+struct PipelineBindContext<'a> {
+    children: &'a [VectorId],
+    batch_inputs: &'a [Vector],
+}
+
+impl BindContext for PipelineBindContext<'_> {
+    fn pipelined_input(&self, pipelined_child_idx: usize) -> VectorId {
+        self.children[pipelined_child_idx]
+    }
+
+    fn batch_input(&self, batch_child_idx: usize) -> Vector {
+        self.batch_inputs[batch_child_idx].clone()
+    }
+}