Zero-copy pipeline input

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

Author: gatesn

encodings/alp/benches/alp_compress.rs

@@ -6,12 +6,12 @@
 use divan::Bencher;
 use rand::rngs::StdRng;
 use rand::{Rng, SeedableRng as _};
-use vortex_alp::{alp_encode, decompress, ALPFloat, ALPRDFloat, RDEncoder};
+use vortex_alp::{ALPFloat, ALPRDFloat, RDEncoder, alp_encode, decompress};
+use vortex_array::IntoArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::compute::warm_up_vtables;
 use vortex_array::validity::Validity;
-use vortex_array::IntoArray;
-use vortex_buffer::{buffer, Buffer};
+use vortex_buffer::{Buffer, buffer};
 use vortex_dtype::NativePType;
 
 fn main() {
@@ -118,7 +118,13 @@ fn decompress_alp_pipeline<T: ALPFloat + NativePType>(bencher: Bencher, args: (u
     bencher
         .with_inputs(|| {
             alp_encode(
-                &PrimitiveArray::new(Buffer::copy_from(&values), validity.clone()),
+                &PrimitiveArray::new(
+                    Buffer::copy_from(&values),
+                    validity
+                        .as_array()
+                        .map(|a| Validity::copy_from_array(a))
+                        .unwrap_or(validity.clone()),
+                ),
                 None,
             )
             .unwrap()

encodings/alp/src/alp/operator.rs

@@ -5,15 +5,15 @@ use std::marker::PhantomData;
 
 use vortex_array::pipeline::bit_view::BitView;
 use vortex_array::pipeline::{
-    BindContext, Kernel, KernelCtx, PipelineInputs, PipelinedNode, VectorId, N,
+    BindContext, Kernel, KernelCtx, N, PipelineInputs, PipelinedNode, VectorId,
 };
 use vortex_array::vtable::OperatorVTable;
 use vortex_dtype::PTypeDowncastExt;
-use vortex_error::{vortex_bail, VortexResult};
+use vortex_error::{VortexResult, vortex_bail};
 use vortex_vector::primitive::PVectorMut;
 use vortex_vector::{VectorMut, VectorMutOps};
 
-use crate::{match_each_alp_float_ptype, ALPArray, ALPFloat, ALPVTable, Exponents};
+use crate::{ALPArray, ALPFloat, ALPVTable, Exponents, match_each_alp_float_ptype};
 
 impl OperatorVTable<ALPVTable> for ALPVTable {
     fn pipeline_node(array: &ALPArray) -> Option<&dyn PipelinedNode> {
@@ -139,11 +139,14 @@ fn sparse_alp<A: ALPFloat>(
 
 #[cfg(test)]
 mod test {
+    use rand::prelude::StdRng;
+    use rand::{Rng, SeedableRng};
+    use vortex_array::IntoArray;
     use vortex_array::arrays::PrimitiveArray;
     use vortex_array::validity::Validity;
-    use vortex_array::IntoArray;
-    use vortex_buffer::buffer;
+    use vortex_buffer::{Buffer, buffer, buffer_mut};
     use vortex_dtype::PTypeDowncastExt;
+    use vortex_vector::VectorOps;
 
     use crate::alp_encode;
 
@@ -161,4 +164,31 @@ mod test {
 
         assert_eq!(decoded.elements(), &buffer);
     }
+
+    #[test]
+    fn test_alp_into_mut() {
+        let (n, fraction_patch, fraction_valid) = (10_000, 0.0, 1.0);
+        let mut rng = StdRng::seed_from_u64(0);
+        let mut values = buffer_mut![1.234f32; n];
+        if fraction_patch > 0.0 {
+            for index in 0..values.len() {
+                if rng.random_bool(fraction_patch) {
+                    values[index] = 1000.0f32
+                }
+            }
+        }
+        let validity = if fraction_valid < 1.0 {
+            Validity::from_iter((0..values.len()).map(|_| rng.random_bool(fraction_valid)))
+        } else {
+            Validity::NonNullable
+        };
+        let values = values.freeze();
+
+        // We take a copy of the values to ensure we can into_mut
+        let array = PrimitiveArray::new(Buffer::copy_from(&values), validity.clone());
+        let array = alp_encode(&array, None).unwrap().into_array();
+
+        let vector = array.execute().unwrap();
+        assert_eq!(vector.len(), values.len());
+    }
 }

vortex-array/src/array/operator.rs

@@ -129,16 +129,17 @@ impl BindCtx for () {
 }
 
 impl dyn Array + '_ {
-    pub fn execute(&self) -> VortexResult<Vector> {
-        self.execute_with_selection(&Mask::new_true(self.len()))
+    pub fn execute(self: Arc<Self>) -> VortexResult<Vector> {
+        let selection = Mask::new_true(self.len());
+        self.execute_with_selection(&selection)
     }
 
-    pub fn execute_with_selection(&self, selection: &Mask) -> VortexResult<Vector> {
+    pub fn execute_with_selection(self: Arc<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);
+            return PipelineDriver::new(self).execute(selection);
         }
 
         self.execute_batch(selection, &mut DummyExecutionCtx)

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

@@ -58,7 +58,7 @@ mod tests {
     fn test_listview_operator_basic() {
         // Test basic ListView execution without selection.
         // ListView: [[0,1,2], [3,4], [5,6], [7,8,9]]
-        let listview = create_basic_listview();
+        let listview = create_basic_listview().into_array();
 
         // Execute without selection.
         let result = listview.execute().unwrap();
@@ -94,7 +94,8 @@ mod tests {
                 .into_array();
         let offsets = PrimitiveArray::from_iter([0u32, 2, 4, 6, 8, 10]).into_array();
         let sizes = PrimitiveArray::from_iter([2u32, 2, 2, 2, 2, 2]).into_array();
-        let listview = ListViewArray::new(elements, offsets, sizes, Validity::AllValid);
+        let listview =
+            ListViewArray::new(elements, offsets, sizes, Validity::AllValid).into_array();
 
         // Create selection mask: [true, false, true, false, true, false].
         let selection = Mask::from_iter([true, false, true, false, true, false]);
@@ -126,7 +127,7 @@ mod tests {
     #[test]
     fn test_listview_operator_with_nulls_and_selection() {
         // Use the nullable listview: [[10,20], null, [50]]
-        let listview = create_nullable_listview();
+        let listview = create_nullable_listview().into_array();
 
         // Create selection mask: [true, true, false].
         let selection = Mask::from_iter([true, true, false]);
@@ -159,7 +160,7 @@ mod tests {
     #[test]
     fn test_listview_operator_overlapping_with_selection() {
         // Use the overlapping listview: [[5,6,7], [2,3], [8,9], [0,1], [1,2,3,4]]
-        let listview = create_overlapping_listview();
+        let listview = create_overlapping_listview().into_array();
 
         // Create selection mask: [true, false, true, true, false].
         let selection = Mask::from_iter([true, false, true, true, false]);

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

@@ -5,34 +5,25 @@ use vortex_buffer::Buffer;
 use vortex_compute::filter::Filter;
 use vortex_dtype::match_each_native_ptype;
 use vortex_error::VortexResult;
+use vortex_mask::Mask;
+use vortex_vector::Vector;
 use vortex_vector::primitive::PVector;
 
 use crate::arrays::{MaskedVTable, PrimitiveArray, PrimitiveVTable};
-use crate::execution::{BatchKernelRef, BindCtx, kernel};
+use crate::execution::ExecutionCtx;
 use crate::vtable::{OperatorVTable, ValidityHelper};
 use crate::{ArrayRef, IntoArray};
 
 impl OperatorVTable<PrimitiveVTable> for PrimitiveVTable {
-    fn bind(
+    fn execute_batch(
         array: &PrimitiveArray,
-        selection: Option<&ArrayRef>,
-        ctx: &mut dyn BindCtx,
-    ) -> VortexResult<BatchKernelRef> {
-        let mask = ctx.bind_selection(array.len(), selection)?;
-        let validity = ctx.bind_validity(array.validity(), array.len(), selection)?;
-
+        selection: &Mask,
+        _ctx: &mut dyn ExecutionCtx,
+    ) -> VortexResult<Vector> {
+        let validity = array.validity_mask().filter(selection);
         match_each_native_ptype!(array.ptype(), |P| {
-            let elements = array.buffer::<P>();
-            Ok(kernel(move || {
-                let mask = mask.execute()?;
-                let validity = validity.execute()?;
-
-                // Note that validity already has the mask applied so we only need to apply it to
-                // the elements.
-                let elements = elements.filter(&mask);
-
-                Ok(PVector::<P>::try_new(elements, validity)?.into())
-            }))
+            let elements = array.buffer::<P>().filter(selection);
+            Ok(PVector::<P>::try_new(elements, validity)?.into())
         })
     }
 

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

@@ -61,7 +61,8 @@ mod tests {
             5,
             Validity::AllValid,
         )
-        .unwrap();
+        .unwrap()
+        .into_array();
 
         // Execute without selection.
         let result = struct_array.execute().unwrap();
@@ -94,7 +95,8 @@ mod tests {
             6,
             Validity::AllValid,
         )
-        .unwrap();
+        .unwrap()
+        .into_array();
 
         // Create a selection mask that selects indices 0, 2, 4 (alternating pattern).
         let selection = Mask::from_iter([true, false, true, false, true, false]);
@@ -146,7 +148,8 @@ mod tests {
             6,
             struct_validity,
         )
-        .unwrap();
+        .unwrap()
+        .into_array();
 
         // Create a selection mask that selects indices 0, 1, 2, 4, 5.
         let selection = Mask::from_iter([true, true, true, false, true, true]);

vortex-array/src/compute/arrays/is_not_null.rs

@@ -128,7 +128,10 @@ mod tests {
         )
         .into_array();
 
-        let result = IsNotNullArray::new(array).execute()?.into_bool();
+        let result = IsNotNullArray::new(array)
+            .into_array()
+            .execute()?
+            .into_bool();
         assert!(result.validity().all_true());
         assert_eq!(result.bits(), &validity);
 

vortex-array/src/compute/arrays/is_null.rs

@@ -131,7 +131,7 @@ mod tests {
         )
         .into_array();
 
-        let result = IsNullArray::new(array).execute()?.into_bool();
+        let result = IsNullArray::new(array).into_array().execute()?.into_bool();
         assert!(result.validity().all_true());
         assert_eq!(result.bits(), &validity.not());
 

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

@@ -6,13 +6,13 @@
 use vortex_error::{VortexExpect, VortexResult};
 use vortex_vector::VectorMut;
 
-use crate::pipeline::driver::{Node, NodeId};
-use crate::pipeline::{VectorId, N};
 use crate::Array;
+use crate::pipeline::driver::{Node, NodeId};
+use crate::pipeline::{N, VectorId};
 
 #[derive(Debug)]
 pub struct VectorAllocation {
-    /// Where each node writes its output
+    /// Where each node writes its output, indexed by [`NodeId`].
     pub(crate) output_targets: Vec<OutputTarget>,
     /// The actual allocated vectors
     pub(crate) vectors: Vec<VectorMut>,
@@ -22,7 +22,7 @@ pub struct VectorAllocation {
 // Node mutates its input in-place (input node index, vector idx)
 // add variant InPlace.
 /// Tracks which vector a node outputs to
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Copy)]
 pub(crate) enum OutputTarget {
     /// Node writes to the top-level provided output
     ExternalOutput,
@@ -52,8 +52,8 @@ pub(super) fn allocate_vectors(
     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];
+    for &node_id in execution_order.iter().rev() {
+        let node = &dag[node_id];
         let array = &node.array;
 
         // Determine output target
@@ -76,7 +76,7 @@ pub(super) fn allocate_vectors(
             OutputTarget::IntermediateVector(vector_id)
         };
 
-        output_targets[node_idx] = Some(output_target);
+        output_targets[node_id] = Some(output_target);
     }
 
     Ok(VectorAllocation {

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

@@ -1,31 +1,41 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::iter;
+
 use vortex_error::{VortexExpect, VortexResult};
 use vortex_vector::{Vector, VectorOps};
 
 use crate::array::ArrayOperator;
 use crate::pipeline::driver::allocation::VectorAllocation;
 use crate::pipeline::driver::input::InputKernel;
-use crate::pipeline::driver::{Node, NodeKind};
+use crate::pipeline::driver::{Node, NodeId, NodeKind};
 use crate::pipeline::{BindContext, Kernel, VectorId};
 
 // We consume the DAG and the batch inputs such that our into_mut calls are safe
 pub(crate) fn bind_kernels(
     dag: Vec<Node>,
+    execution_order: &[NodeId],
     allocation_plan: &VectorAllocation,
     mut all_batch_inputs: Vec<Option<Vector>>,
 ) -> VortexResult<Vec<Box<dyn Kernel>>> {
-    let mut kernels = Vec::with_capacity(dag.len());
-    for node in dag {
+    // We construct kernels in top-down order (i.e. inverse execution order) such that any arrays
+    // in parent nodes are dropped prior to constructing their children. This gives us the best
+    // chance of performing zero-copy into_mut calls on the pipeline input arrays.
+
+    let mut dag: Vec<_> = dag.into_iter().map(Some).collect();
+    let mut kernels: Vec<Option<Box<dyn Kernel>>> =
+        iter::repeat_with(|| None).take(dag.len()).collect();
+
+    // Note that the execution order is bottom-up, so to make sure we consume an array's parent
+    // before construct the array's own kernel, we must construct kernels in reverse order.
+    for node_id in execution_order.iter().copied().rev() {
+        let node = dag[node_id].take().vortex_expect("already processed");
+
         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")
-            })
+            .map(|child_node_id| allocation_plan.output_targets[*child_node_id].vector_id())
             .collect::<Vec<_>>();
 
         let mut batch_inputs: Vec<_> = node
@@ -34,7 +44,7 @@ pub(crate) fn bind_kernels(
             .map(|idx| all_batch_inputs[*idx].take())
             .collect();
 
-        kernels.push(match node.array.as_pipelined() {
+        kernels[node_id] = Some(match node.array.as_pipelined() {
             None => {
                 // If the node cannot be pipelined, it must be an input node
                 assert_eq!(node.kind, NodeKind::Input);
@@ -61,17 +71,22 @@ pub(crate) fn bind_kernels(
             }
         });
     }
-    Ok(kernels)
+
+    Ok(kernels
+        .into_iter()
+        .map(|k| k.vortex_expect("missing kernel"))
+        .collect())
 }
 
 struct PipelineBindContext<'a> {
-    children: &'a [VectorId],
+    children: &'a [Option<VectorId>],
     batch_inputs: &'a mut [Option<Vector>],
 }
 
 impl BindContext for PipelineBindContext<'_> {
     fn pipelined_input(&self, pipelined_child_idx: usize) -> VectorId {
         self.children[pipelined_child_idx]
+            .vortex_expect("In-place transforms do not have VectorIDs for their zero'th child.")
     }
 
     fn batch_input(&mut self, batch_child_idx: usize) -> Vector {