Use cursor copies for PiecewiseSequence slices

Signed-off-by: Daniel King <dan@spiraldb.com>

Author: danking

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

@@ -17,6 +17,7 @@ use crate::arrays::PiecewiseSequence;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::dict::TakeExecute;
 use crate::arrays::piecewise_sequence::ConstantOrArray;
+use crate::arrays::piecewise_sequence::copy_slice_to_spare;
 use crate::arrays::piecewise_sequence::maybe_contiguous_slices;
 use crate::dtype::IntegerPType;
 use crate::dtype::NativeDecimalType;
@@ -196,11 +197,19 @@ where
     );
 
     let mut result = BufferMut::<T>::with_capacity(output_len);
+    let mut cursor = 0usize;
     for &start in starts {
         let start = start.as_();
-        result.extend_from_slice(&values[start..][..length]);
+        cursor = copy_slice_to_spare(&mut result, cursor, &values[start..][..length], output_len)?;
     }
 
+    vortex_ensure!(
+        cursor == output_len,
+        "PiecewiseSequenceArray expanded length {cursor} does not match declared length {output_len}"
+    );
+    // SAFETY: `copy_slice_to_spare` checked every write against `output_len`, and `cursor ==
+    // output_len` proves all slots were initialized.
+    unsafe { result.set_len(output_len) };
     Ok(result.freeze())
 }
 
@@ -216,17 +225,21 @@ where
     T: NativeDecimalType,
 {
     let mut result = BufferMut::<T>::with_capacity(output_len);
+    let mut cursor = 0usize;
     for (&start, &length) in starts.iter().zip_eq(lengths) {
         let start = start.as_();
         let length = length.as_();
-        result.extend_from_slice(&values[start..][..length]);
+        cursor = copy_slice_to_spare(&mut result, cursor, &values[start..][..length], output_len)?;
     }
 
     vortex_ensure!(
-        result.len() == output_len,
+        cursor == output_len,
         "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}",
-        result.len()
+        cursor
     );
+    // SAFETY: `copy_slice_to_spare` checked every write against `output_len`, and `cursor ==
+    // output_len` proves all slots were initialized.
+    unsafe { result.set_len(output_len) };
     Ok(result.freeze())
 }
 

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

@@ -8,6 +8,8 @@
 //! intended for take operations that can gather regular runs without materializing one index per
 //! element.
 
+use std::ptr;
+
 use itertools::Itertools;
 use num_traits::AsPrimitive;
 use vortex_buffer::BufferMut;
@@ -102,6 +104,40 @@ pub(crate) fn is_constant_multiplier_one(multipliers: &ArrayRef) -> bool {
     )
 }
 
+pub(crate) fn copy_slice_to_spare<T: Copy>(
+    buffer: &mut BufferMut<T>,
+    cursor: usize,
+    source: &[T],
+    output_len: usize,
+) -> VortexResult<usize> {
+    let end = cursor
+        .checked_add(source.len())
+        .ok_or_else(|| vortex_err!("slice copy output length overflows usize"))?;
+    vortex_ensure!(
+        end <= output_len,
+        "slice copy length {end} exceeds declared output length {output_len}"
+    );
+    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()
+    );
+
+    // SAFETY: the checks above prove `cursor..end` is inside the spare capacity of an
+    // uninitialized buffer allocated for at least `output_len` values.
+    let dst = unsafe { buffer.spare_capacity_mut().get_unchecked_mut(cursor..end) };
+    // SAFETY: `dst` has exactly `source.len()` writable slots and does not overlap with `source`.
+    unsafe {
+        ptr::copy_nonoverlapping(source.as_ptr(), dst.as_mut_ptr().cast(), source.len());
+    }
+    Ok(end)
+}
+
 fn constant_unsigned_usize(array: &ArrayRef) -> VortexResult<Option<usize>> {
     let Some(scalar) = array.as_constant() else {
         return Ok(None);

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

@@ -24,6 +24,7 @@ use crate::arrays::Primitive;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::dict::TakeExecute;
 use crate::arrays::piecewise_sequence::ConstantOrArray;
+use crate::arrays::piecewise_sequence::copy_slice_to_spare;
 use crate::arrays::piecewise_sequence::maybe_contiguous_slices;
 use crate::builtins::ArrayBuiltins;
 use crate::dtype::DType;
@@ -242,17 +243,21 @@ where
     L: UnsignedPType,
 {
     let mut values = BufferMut::<T>::with_capacity(output_len);
+    let mut cursor = 0usize;
     for (&start, &length) in starts.iter().zip_eq(lengths) {
         let start = start.as_();
         let length = length.as_();
-        values.extend_from_slice(&source[start..][..length]);
+        cursor = copy_slice_to_spare(&mut values, cursor, &source[start..][..length], output_len)?;
     }
 
     vortex_ensure!(
-        values.len() == output_len,
+        cursor == output_len,
         "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}",
-        values.len()
+        cursor
     );
+    // SAFETY: `copy_slice_to_spare` checked every write against `output_len`, and `cursor ==
+    // output_len` proves all slots were initialized.
+    unsafe { values.set_len(output_len) };
     Ok(values.freeze())
 }
 
@@ -309,11 +314,19 @@ where
     );
 
     let mut values = BufferMut::<T>::with_capacity(output_len);
+    let mut cursor = 0usize;
     for &start in starts {
         let start = start.as_();
-        values.extend_from_slice(&source[start..][..length]);
+        cursor = copy_slice_to_spare(&mut values, cursor, &source[start..][..length], output_len)?;
     }
 
+    vortex_ensure!(
+        cursor == output_len,
+        "PiecewiseSequenceArray expanded length {cursor} does not match declared length {output_len}"
+    );
+    // SAFETY: `copy_slice_to_spare` checked every write against `output_len`, and `cursor ==
+    // output_len` proves all slots were initialized.
+    unsafe { values.set_len(output_len) };
     Ok(values.freeze())
 }
 

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

@@ -21,6 +21,7 @@ use crate::arrays::VarBin;
 use crate::arrays::VarBinArray;
 use crate::arrays::dict::TakeExecute;
 use crate::arrays::piecewise_sequence::ConstantOrArray;
+use crate::arrays::piecewise_sequence::copy_slice_to_spare;
 use crate::arrays::piecewise_sequence::maybe_contiguous_slices;
 use crate::arrays::primitive::PrimitiveArrayExt;
 use crate::arrays::varbin::VarBinArrayExt;
@@ -467,6 +468,7 @@ where
     }
 
     let mut new_data = ByteBufferMut::with_capacity(output_bytes);
+    let mut cursor = 0usize;
     for &start in starts {
         let start = start.as_();
         if length == 0 {
@@ -476,8 +478,20 @@ 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]);
+        cursor = copy_slice_to_spare(
+            &mut new_data,
+            cursor,
+            &data[byte_start..][..byte_end - byte_start],
+            output_bytes,
+        )?;
     }
+    vortex_ensure!(
+        cursor == output_bytes,
+        "VarBin byte copy length {cursor} does not match computed byte length {output_bytes}"
+    );
+    // SAFETY: `copy_slice_to_spare` checked every write against `output_bytes`, and `cursor ==
+    // output_bytes` proves all bytes were initialized.
+    unsafe { new_data.set_len(output_bytes) };
 
     let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable)
         .reinterpret_cast(out_offset_ptype)
@@ -544,6 +558,7 @@ where
     );
 
     let mut new_data = ByteBufferMut::with_capacity(output_bytes);
+    let mut cursor = 0usize;
     for (&start, &length) in starts.iter().zip_eq(lengths) {
         let start = start.as_();
         let length = length.as_();
@@ -554,8 +569,20 @@ 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]);
+        cursor = copy_slice_to_spare(
+            &mut new_data,
+            cursor,
+            &data[byte_start..][..byte_end - byte_start],
+            output_bytes,
+        )?;
     }
+    vortex_ensure!(
+        cursor == output_bytes,
+        "VarBin byte copy length {cursor} does not match computed byte length {output_bytes}"
+    );
+    // SAFETY: `copy_slice_to_spare` checked every write against `output_bytes`, and `cursor ==
+    // output_bytes` proves all bytes were initialized.
+    unsafe { new_data.set_len(output_bytes) };
 
     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::VarBinView;
 use crate::arrays::VarBinViewArray;
 use crate::arrays::dict::TakeExecute;
 use crate::arrays::piecewise_sequence::ConstantOrArray;
+use crate::arrays::piecewise_sequence::copy_slice_to_spare;
 use crate::arrays::piecewise_sequence::maybe_contiguous_slices;
 use crate::arrays::varbinview::BinaryView;
 use crate::buffer::BufferHandle;
@@ -171,11 +172,19 @@ where
     );
 
     let mut views = BufferMut::<BinaryView>::with_capacity(output_len);
+    let mut cursor = 0usize;
     for &start in starts {
         let start = start.as_();
-        views.extend_from_slice(&source[start..][..length]);
+        cursor = copy_slice_to_spare(&mut views, cursor, &source[start..][..length], output_len)?;
     }
 
+    vortex_ensure!(
+        cursor == output_len,
+        "PiecewiseSequenceArray expanded length {cursor} does not match declared length {output_len}"
+    );
+    // SAFETY: `copy_slice_to_spare` checked every write against `output_len`, and `cursor ==
+    // output_len` proves all slots were initialized.
+    unsafe { views.set_len(output_len) };
     Ok(views.freeze())
 }
 
@@ -190,17 +199,21 @@ where
     L: UnsignedPType,
 {
     let mut views = BufferMut::<BinaryView>::with_capacity(output_len);
+    let mut cursor = 0usize;
     for (&start, &length) in starts.iter().zip_eq(lengths) {
         let start = start.as_();
         let length = length.as_();
-        views.extend_from_slice(&source[start..][..length]);
+        cursor = copy_slice_to_spare(&mut views, cursor, &source[start..][..length], output_len)?;
     }
 
     vortex_ensure!(
-        views.len() == output_len,
+        cursor == output_len,
         "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}",
-        views.len()
+        cursor
     );
+    // SAFETY: `copy_slice_to_spare` checked every write against `output_len`, and `cursor ==
+    // output_len` proves all slots were initialized.
+    unsafe { views.set_len(output_len) };
     Ok(views.freeze())
 }