Use cursor copies for PiecewiseSequence slices

vortex-array/src/arrays/decimal/compute/take.rs

@@ -17,6 +17,7 @@ use crate::arrays::DecimalArray;
 use crate::arrays::PiecewiseSequence;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::dict::TakeExecute;
+use crate::arrays::piecewise_sequence::SpareBufferWriter;
 use crate::arrays::piecewise_sequence::constant_unsigned_usize;
 use crate::arrays::piecewise_sequence::maybe_contiguous_slices;
 use crate::dtype::IntegerPType;
@@ -199,11 +200,12 @@ where
     );
 
     let mut result = BufferMut::<T>::with_capacity(output_len);
+    let mut writer = SpareBufferWriter::new(&mut result, output_len)?;
     for &start in starts {
         let start = start.as_();
-        result.extend_from_slice(&values[start..][..length]);
+        writer.copy_slice(&values[start..][..length])?;
     }
-
+    writer.finish()?;
     Ok(result.freeze())
 }
 
@@ -219,17 +221,13 @@ where
     T: NativeDecimalType,
 {
     let mut result = BufferMut::<T>::with_capacity(output_len);
+    let mut writer = SpareBufferWriter::new(&mut result, output_len)?;
     for (&start, &length) in starts.iter().zip_eq(lengths) {
         let start = start.as_();
         let length = length.as_();
-        result.extend_from_slice(&values[start..][..length]);
+        writer.copy_slice(&values[start..][..length])?;
     }
-
-    vortex_ensure!(
-        result.len() == output_len,
-        "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}",
-        result.len()
-    );
+    writer.finish()?;
     Ok(result.freeze())
 }
 

vortex-array/src/arrays/piecewise_sequence/mod.rs

@@ -28,12 +28,14 @@ use crate::executor::ExecutionCtx;
 use crate::scalar::PValue;
 
 pub mod array;
+mod spare_buffer_writer;
 mod vtable;
 
 #[cfg(test)]
 mod tests;
 
 pub use array::PiecewiseSequenceArrayExt;
+pub(crate) use spare_buffer_writer::SpareBufferWriter;
 pub use vtable::*;
 
 pub(crate) fn check_index_arrays(

vortex-array/src/arrays/piecewise_sequence/spare_buffer_writer.rs

@@ -0,0 +1,140 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::ptr;
+
+use vortex_buffer::BufferMut;
+use vortex_error::VortexResult;
+use vortex_error::vortex_ensure;
+
+/// Writes slices sequentially into the spare capacity of an empty [`BufferMut`].
+///
+/// This is useful when the caller knows the final output length up front and wants to avoid
+/// repeatedly growing the initialized length while copying contiguous source slices.
+///
+/// All writes are bounds-checked against the declared output length; the check is a single
+/// well-predicted branch per slice and is dominated by the copy itself. The unsafety is
+/// contained to the copy into the bounds-checked spare-capacity range and the final
+/// [`set_len`], which [`finish`] justifies from the writer's invariant that copies initialize
+/// a contiguous prefix of the spare capacity.
+///
+/// [`set_len`]: BufferMut::set_len
+/// [`finish`]: SpareBufferWriter::finish
+#[must_use = "call `finish` to set the buffer length after writing"]
+pub(crate) struct SpareBufferWriter<'a, T> {
+    buffer: &'a mut BufferMut<T>,
+    written: usize,
+    output_len: usize,
+}
+
+impl<'a, T: Copy> SpareBufferWriter<'a, T> {
+    /// Creates a writer for `output_len` values in `buffer`'s spare capacity.
+    ///
+    /// The target buffer must be empty and have capacity for at least `output_len` values.
+    pub(crate) fn new(buffer: &'a mut BufferMut<T>, output_len: usize) -> VortexResult<Self> {
+        vortex_ensure!(
+            buffer.is_empty(),
+            "slice copy buffer already has {} initialized values",
+            buffer.len()
+        );
+        vortex_ensure!(
+            output_len <= buffer.capacity(),
+            "slice copy output length {output_len} exceeds buffer capacity {}",
+            buffer.capacity()
+        );
+
+        Ok(Self {
+            buffer,
+            written: 0,
+            output_len,
+        })
+    }
+
+    /// Copies `source` into the next output slots.
+    #[inline]
+    pub(crate) fn copy_slice(&mut self, source: &[T]) -> VortexResult<()> {
+        vortex_ensure!(
+            source.len() <= self.output_len - self.written,
+            "slice copy length {} exceeds remaining output length {}",
+            source.len(),
+            self.output_len - self.written
+        );
+        let end = self.written + source.len();
+        let dst = &mut self.buffer.spare_capacity_mut()[self.written..end];
+        // SAFETY: `MaybeUninit<T>` has the same layout as `T`, `dst` and `source` have equal
+        // lengths, and `dst` lives in the exclusively borrowed buffer so the two cannot overlap.
+        // This is `<[MaybeUninit<T>]>::write_copy_of_slice`, which is not yet stable.
+        unsafe {
+            ptr::copy_nonoverlapping(source.as_ptr(), dst.as_mut_ptr().cast::<T>(), source.len());
+        }
+        self.written = end;
+        Ok(())
+    }
+
+    /// Sets the target buffer length after exactly `output_len` values have been written.
+    pub(crate) fn finish(self) -> VortexResult<()> {
+        vortex_ensure!(
+            self.written == self.output_len,
+            "slice copy length {} does not match declared output length {}",
+            self.written,
+            self.output_len
+        );
+        // SAFETY: `copy_slice` initializes the contiguous prefix `0..written` of the spare
+        // capacity, and `new` checked that `output_len` does not exceed the capacity.
+        unsafe {
+            self.buffer.set_len(self.output_len);
+        }
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_buffer::BufferMut;
+    use vortex_error::VortexResult;
+
+    use super::SpareBufferWriter;
+
+    #[test]
+    fn writes_slices_into_spare_capacity() -> VortexResult<()> {
+        let mut buffer = BufferMut::with_capacity(4);
+
+        let mut writer = SpareBufferWriter::new(&mut buffer, 4)?;
+        writer.copy_slice(&[1u16, 2])?;
+        writer.copy_slice(&[3u16, 4])?;
+        writer.finish()?;
+
+        assert_eq!(buffer.as_slice(), &[1u16, 2, 3, 4]);
+        Ok(())
+    }
+
+    #[test]
+    fn rejects_overlong_copy() -> VortexResult<()> {
+        let mut buffer = BufferMut::with_capacity(2);
+
+        let mut writer = SpareBufferWriter::new(&mut buffer, 2)?;
+        let error = writer.copy_slice(&[1u16, 2, 3]).unwrap_err();
+
+        assert!(
+            error
+                .to_string()
+                .contains("exceeds remaining output length")
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn rejects_incomplete_finish() -> VortexResult<()> {
+        let mut buffer = BufferMut::with_capacity(2);
+
+        let writer = SpareBufferWriter::<u16>::new(&mut buffer, 2)?;
+        let error = writer.finish().unwrap_err();
+
+        assert!(
+            error
+                .to_string()
+                .contains("does not match declared output length")
+        );
+        Ok(())
+    }
+}

vortex-array/src/arrays/primitive/compute/take/mod.rs

@@ -24,6 +24,7 @@ use crate::arrays::PiecewiseSequence;
 use crate::arrays::Primitive;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::dict::TakeExecute;
+use crate::arrays::piecewise_sequence::SpareBufferWriter;
 use crate::arrays::piecewise_sequence::constant_unsigned_usize;
 use crate::arrays::piecewise_sequence::maybe_contiguous_slices;
 use crate::builtins::ArrayBuiltins;
@@ -245,17 +246,13 @@ where
     L: UnsignedPType,
 {
     let mut values = BufferMut::<T>::with_capacity(output_len);
+    let mut writer = SpareBufferWriter::new(&mut values, output_len)?;
     for (&start, &length) in starts.iter().zip_eq(lengths) {
         let start = start.as_();
         let length = length.as_();
-        values.extend_from_slice(&source[start..][..length]);
+        writer.copy_slice(&source[start..][..length])?;
     }
-
-    vortex_ensure!(
-        values.len() == output_len,
-        "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}",
-        values.len()
-    );
+    writer.finish()?;
     Ok(values.freeze())
 }
 
@@ -312,11 +309,12 @@ where
     );
 
     let mut values = BufferMut::<T>::with_capacity(output_len);
+    let mut writer = SpareBufferWriter::new(&mut values, output_len)?;
     for &start in starts {
         let start = start.as_();
-        values.extend_from_slice(&source[start..][..length]);
+        writer.copy_slice(&source[start..][..length])?;
     }
-
+    writer.finish()?;
     Ok(values.freeze())
 }
 

vortex-array/src/arrays/varbin/compute/take.rs

@@ -21,6 +21,7 @@ use crate::arrays::PrimitiveArray;
 use crate::arrays::VarBin;
 use crate::arrays::VarBinArray;
 use crate::arrays::dict::TakeExecute;
+use crate::arrays::piecewise_sequence::SpareBufferWriter;
 use crate::arrays::piecewise_sequence::constant_unsigned_usize;
 use crate::arrays::piecewise_sequence::maybe_contiguous_slices;
 use crate::arrays::primitive::PrimitiveArrayExt;
@@ -474,6 +475,7 @@ where
     }
 
     let mut new_data = ByteBufferMut::with_capacity(output_bytes);
+    let mut writer = SpareBufferWriter::new(&mut new_data, output_bytes)?;
     for &start in starts {
         let start = start.as_();
         if length == 0 {
@@ -483,8 +485,9 @@ where
         let offset_range = &offsets[start..][..=length];
         let byte_start = offset_range[0].as_();
         let byte_end = offset_range[length].as_();
-        new_data.extend_from_slice(&data[byte_start..][..byte_end - byte_start]);
+        writer.copy_slice(&data[byte_start..][..byte_end - byte_start])?;
     }
+    writer.finish()?;
 
     let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable)
         .reinterpret_cast(out_offset_ptype)
@@ -551,6 +554,7 @@ where
     );
 
     let mut new_data = ByteBufferMut::with_capacity(output_bytes);
+    let mut writer = SpareBufferWriter::new(&mut new_data, output_bytes)?;
     for (&start, &length) in starts.iter().zip_eq(lengths) {
         let start = start.as_();
         let length = length.as_();
@@ -561,8 +565,9 @@ where
         let offset_range = &offsets[start..][..=length];
         let byte_start = offset_range[0].as_();
         let byte_end = offset_range[length].as_();
-        new_data.extend_from_slice(&data[byte_start..][..byte_end - byte_start]);
+        writer.copy_slice(&data[byte_start..][..byte_end - byte_start])?;
     }
+    writer.finish()?;
 
     let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable)
         .reinterpret_cast(out_offset_ptype)

vortex-array/src/arrays/varbinview/compute/take.rs

@@ -23,6 +23,7 @@ use crate::arrays::PrimitiveArray;
 use crate::arrays::VarBinView;
 use crate::arrays::VarBinViewArray;
 use crate::arrays::dict::TakeExecute;
+use crate::arrays::piecewise_sequence::SpareBufferWriter;
 use crate::arrays::piecewise_sequence::constant_unsigned_usize;
 use crate::arrays::piecewise_sequence::maybe_contiguous_slices;
 use crate::arrays::varbinview::BinaryView;
@@ -174,11 +175,12 @@ where
     );
 
     let mut views = BufferMut::<BinaryView>::with_capacity(output_len);
+    let mut writer = SpareBufferWriter::new(&mut views, output_len)?;
     for &start in starts {
         let start = start.as_();
-        views.extend_from_slice(&source[start..][..length]);
+        writer.copy_slice(&source[start..][..length])?;
     }
-
+    writer.finish()?;
     Ok(views.freeze())
 }
 
@@ -193,17 +195,13 @@ where
     L: UnsignedPType,
 {
     let mut views = BufferMut::<BinaryView>::with_capacity(output_len);
+    let mut writer = SpareBufferWriter::new(&mut views, output_len)?;
     for (&start, &length) in starts.iter().zip_eq(lengths) {
         let start = start.as_();
         let length = length.as_();
-        views.extend_from_slice(&source[start..][..length]);
+        writer.copy_slice(&source[start..][..length])?;
     }
-
-    vortex_ensure!(
-        views.len() == output_len,
-        "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}",
-        views.len()
-    );
+    writer.finish()?;
     Ok(views.freeze())
 }