Fix zero-copy pipeline input (#5324)

The into_mut was triggering a copy before this

---------

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

Author: gatesn

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

@@ -11,7 +11,7 @@ use crate::pipeline::driver::{Node, NodeKind};
 use crate::pipeline::{BindContext, Kernel, VectorId};
 
 pub(crate) fn bind_kernels(
-    dag: &[Node],
+    dag: Vec<Node>,
     allocation_plan: &VectorAllocation,
     mut all_batch_inputs: Vec<Option<Vector>>,
 ) -> VortexResult<Vec<Box<dyn Kernel>>> {
@@ -40,6 +40,10 @@ pub(crate) fn bind_kernels(
                 assert_eq!(node.batch_inputs.len(), 1);
                 let batch_id = node.batch_inputs[0];
 
+                // Release ownership of the array before trying to call into_mut on the vector.
+                // This is in case the vector was constructed zero-copy from the array's data.
+                drop(node.array);
+
                 let batch = batch_inputs[batch_id]
                     .take()
                     .vortex_expect("Batch input vector has already been consumed")

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

@@ -73,11 +73,13 @@ impl Kernel for InputKernel {
 
 #[cfg(test)]
 mod test {
-    use vortex_buffer::{bitbuffer, buffer};
+    use vortex_buffer::{BitBuffer, bitbuffer, buffer};
     use vortex_dtype::PTypeDowncastExt;
     use vortex_mask::Mask;
 
+    use crate::arrays::{BoolArray, PrimitiveArray};
     use crate::pipeline::driver::PipelineDriver;
+    use crate::validity::Validity;
     use crate::{Array, ArrayOperator, IntoArray};
 
     #[test]
@@ -113,4 +115,38 @@ mod test {
             .downcast::<u32>();
         assert_eq!(vector.elements().as_ref(), &[0u32, 2, 4]);
     }
+
+    /// Ensures that we can feed an input into a pipeline with zero-copy.
+    /// This can require careful book keeping to make sure we don't hold references to arrays
+    /// around longer than necessary.
+    #[test]
+    fn test_pipeline_input_zero_copy() {
+        let elements = buffer![123u32; 8000];
+        let elements_ptr = elements.as_ptr();
+        let validity = BitBuffer::from_iter((0..8000).map(|i| i % 2 == 0));
+        let validity_ptr = validity.inner().as_ptr();
+
+        let array = PrimitiveArray::new(
+            elements,
+            Validity::Array(BoolArray::from(validity).into_array()),
+        )
+        .into_array();
+        assert!(
+            array.as_pipelined().is_none(),
+            "We're explicitly testing non-pipelined arrays to trigger the input case"
+        );
+
+        let selection = Mask::new_true(array.len());
+        let vector = PipelineDriver::new(array)
+            .execute(&selection)
+            .unwrap()
+            .into_primitive()
+            .downcast::<u32>();
+
+        let (vector_elements, vector_validity) = vector.into_parts();
+        let vector_validity = vector_validity.into_bit_buffer().into_inner();
+
+        assert_eq!(vector_elements.as_ptr(), elements_ptr);
+        assert_eq!(vector_validity.as_ptr(), validity_ptr);
+    }
 }

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

@@ -209,7 +209,7 @@ impl PipelineDriver {
         let allocation_plan = allocate_vectors(&self.dag, &exec_order)?;
 
         // Bind each node in the DAG to create its kernel
-        let kernels = bind_kernels(&self.dag, &allocation_plan, batch_inputs)?;
+        let kernels = bind_kernels(self.dag, &allocation_plan, batch_inputs)?;
 
         // Construct the kernel execution context
         let ctx = KernelCtx::new(allocation_plan.vectors);

vortex-mask/src/lib.rs

@@ -441,6 +441,19 @@ impl Mask {
         }
     }
 
+    /// Return a boolean buffer representation of the mask, allocating new buffers for all-true
+    /// and all-false variants.
+    #[inline]
+    pub fn into_bit_buffer(self) -> BitBuffer {
+        match self {
+            Self::AllTrue(l) => BitBuffer::new_set(l),
+            Self::AllFalse(l) => BitBuffer::new_unset(l),
+            Self::Values(values) => Arc::try_unwrap(values)
+                .map(|v| v.into_bit_buffer())
+                .unwrap_or_else(|v| v.bit_buffer().clone()),
+        }
+    }
+
     /// Return the indices representation of the mask.
     #[inline]
     pub fn indices(&self) -> AllOr<&[usize]> {
@@ -598,6 +611,12 @@ impl MaskValues {
         &self.buffer
     }
 
+    /// Returns the boolean buffer representation of the mask.
+    #[inline]
+    pub fn into_bit_buffer(self) -> BitBuffer {
+        self.buffer
+    }
+
     /// Returns the boolean value at a given index.
     #[inline]
     pub fn value(&self, index: usize) -> bool {