Advance PiecewiseSequence copy destinations

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

Author: danking

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

@@ -17,8 +17,8 @@ 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::copy_slice_to_spare;
 use crate::arrays::piecewise_sequence::maybe_contiguous_slices;
 use crate::dtype::IntegerPType;
 use crate::dtype::NativeDecimalType;
@@ -200,19 +200,14 @@ where
     );
 
     let mut result = BufferMut::<T>::with_capacity(output_len);
-    let mut cursor = 0usize;
+    let mut writer = SpareBufferWriter::new(&mut result, output_len)?;
     for &start in starts {
         let start = start.as_();
-        cursor = copy_slice_to_spare(&mut result, cursor, &values[start..][..length], output_len)?;
+        // SAFETY: `computed_len == output_len` proves that all fixed-length slices fit in the
+        // output buffer.
+        unsafe { writer.copy_slice_unchecked(&values[start..][..length]) };
     }
-
-    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) };
+    writer.finish()?;
     Ok(result.freeze())
 }
 
@@ -228,21 +223,13 @@ where
     T: NativeDecimalType,
 {
     let mut result = BufferMut::<T>::with_capacity(output_len);
-    let mut cursor = 0usize;
+    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_();
-        cursor = copy_slice_to_spare(&mut result, cursor, &values[start..][..length], output_len)?;
+        writer.copy_slice(&values[start..][..length])?;
     }
-
-    vortex_ensure!(
-        cursor == output_len,
-        "PiecewiseSequenceArray expanded length {} does not match declared length {output_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) };
+    writer.finish()?;
     Ok(result.freeze())
 }
 

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

@@ -94,38 +94,78 @@ pub(crate) fn is_constant_one(multipliers: &ArrayRef) -> bool {
     )
 }
 
-pub(crate) fn copy_slice_to_spare<T: Copy>(
-    buffer: &mut BufferMut<T>,
-    cursor: usize,
-    source: &[T],
+pub(crate) struct SpareBufferWriter<'a, T> {
+    buffer: &'a mut BufferMut<T>,
+    next: *mut T,
+    remaining: usize,
     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()
-    );
+}
+
+impl<'a, T: Copy> SpareBufferWriter<'a, T> {
+    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()
+        );
 
-    // 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());
+        let next = buffer.spare_capacity_mut().as_mut_ptr().cast();
+        Ok(Self {
+            buffer,
+            next,
+            remaining: output_len,
+            output_len,
+        })
+    }
+
+    #[inline]
+    pub(crate) fn copy_slice(&mut self, source: &[T]) -> VortexResult<()> {
+        vortex_ensure!(
+            source.len() <= self.remaining,
+            "slice copy length {} exceeds remaining output length {}",
+            source.len(),
+            self.remaining
+        );
+        // SAFETY: the check above proves that `source` fits in the remaining output slots.
+        unsafe { self.copy_slice_unchecked(source) };
+        Ok(())
+    }
+
+    /// Copies `source` to the next output slots without checking the remaining output length.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that `source.len()` does not exceed the remaining output length.
+    #[inline]
+    pub(crate) unsafe fn copy_slice_unchecked(&mut self, source: &[T]) {
+        debug_assert!(source.len() <= self.remaining);
+        // SAFETY: `next` points to the first unwritten slot, the caller guarantees the source fits
+        // in the remaining capacity, and the mutable buffer borrow prevents reallocation.
+        unsafe {
+            ptr::copy_nonoverlapping(source.as_ptr(), self.next, source.len());
+            self.next = self.next.add(source.len());
+        }
+        self.remaining -= source.len();
+    }
+
+    pub(crate) fn finish(self) -> VortexResult<()> {
+        vortex_ensure!(
+            self.remaining == 0,
+            "slice copy length {} does not match declared output length {}",
+            self.output_len - self.remaining,
+            self.output_len
+        );
+        // SAFETY: successful calls to the copy methods initialized exactly `output_len` slots.
+        unsafe {
+            self.buffer.set_len(self.output_len);
+        }
+        Ok(())
     }
-    Ok(end)
 }
 
 pub(crate) fn constant_unsigned_usize(array: &ConstantArray) -> usize {

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

@@ -24,8 +24,8 @@ 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::copy_slice_to_spare;
 use crate::arrays::piecewise_sequence::maybe_contiguous_slices;
 use crate::builtins::ArrayBuiltins;
 use crate::dtype::DType;
@@ -246,21 +246,13 @@ where
     L: UnsignedPType,
 {
     let mut values = BufferMut::<T>::with_capacity(output_len);
-    let mut cursor = 0usize;
+    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_();
-        cursor = copy_slice_to_spare(&mut values, cursor, &source[start..][..length], output_len)?;
+        writer.copy_slice(&source[start..][..length])?;
     }
-
-    vortex_ensure!(
-        cursor == output_len,
-        "PiecewiseSequenceArray expanded length {} does not match declared length {output_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) };
+    writer.finish()?;
     Ok(values.freeze())
 }
 
@@ -317,19 +309,14 @@ where
     );
 
     let mut values = BufferMut::<T>::with_capacity(output_len);
-    let mut cursor = 0usize;
+    let mut writer = SpareBufferWriter::new(&mut values, output_len)?;
     for &start in starts {
         let start = start.as_();
-        cursor = copy_slice_to_spare(&mut values, cursor, &source[start..][..length], output_len)?;
+        // SAFETY: `computed_len == output_len` proves that all fixed-length slices fit in the
+        // output buffer.
+        unsafe { writer.copy_slice_unchecked(&source[start..][..length]) };
     }
-
-    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) };
+    writer.finish()?;
     Ok(values.freeze())
 }
 

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

@@ -21,8 +21,8 @@ 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::copy_slice_to_spare;
 use crate::arrays::piecewise_sequence::maybe_contiguous_slices;
 use crate::arrays::primitive::PrimitiveArrayExt;
 use crate::arrays::varbin::VarBinArrayExt;
@@ -475,7 +475,7 @@ where
     }
 
     let mut new_data = ByteBufferMut::with_capacity(output_bytes);
-    let mut cursor = 0usize;
+    let mut writer = SpareBufferWriter::new(&mut new_data, output_bytes)?;
     for &start in starts {
         let start = start.as_();
         if length == 0 {
@@ -485,20 +485,10 @@ where
         let offset_range = &offsets[start..][..=length];
         let byte_start = offset_range[0].as_();
         let byte_end = offset_range[length].as_();
-        cursor = copy_slice_to_spare(
-            &mut new_data,
-            cursor,
-            &data[byte_start..][..byte_end - byte_start],
-            output_bytes,
-        )?;
+        // SAFETY: the first pass computed `output_bytes` from these same byte ranges.
+        unsafe { writer.copy_slice_unchecked(&data[byte_start..][..byte_end - byte_start]) };
     }
-    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) };
+    writer.finish()?;
 
     let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable)
         .reinterpret_cast(out_offset_ptype)
@@ -565,7 +555,7 @@ where
     );
 
     let mut new_data = ByteBufferMut::with_capacity(output_bytes);
-    let mut cursor = 0usize;
+    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_();
@@ -576,20 +566,10 @@ where
         let offset_range = &offsets[start..][..=length];
         let byte_start = offset_range[0].as_();
         let byte_end = offset_range[length].as_();
-        cursor = copy_slice_to_spare(
-            &mut new_data,
-            cursor,
-            &data[byte_start..][..byte_end - byte_start],
-            output_bytes,
-        )?;
+        // SAFETY: the first pass computed `output_bytes` from these same byte ranges.
+        unsafe { writer.copy_slice_unchecked(&data[byte_start..][..byte_end - byte_start]) };
     }
-    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) };
+    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,8 +23,8 @@ 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::copy_slice_to_spare;
 use crate::arrays::piecewise_sequence::maybe_contiguous_slices;
 use crate::arrays::varbinview::BinaryView;
 use crate::buffer::BufferHandle;
@@ -175,19 +175,14 @@ where
     );
 
     let mut views = BufferMut::<BinaryView>::with_capacity(output_len);
-    let mut cursor = 0usize;
+    let mut writer = SpareBufferWriter::new(&mut views, output_len)?;
     for &start in starts {
         let start = start.as_();
-        cursor = copy_slice_to_spare(&mut views, cursor, &source[start..][..length], output_len)?;
+        // SAFETY: `computed_len == output_len` proves that all fixed-length slices fit in the
+        // output buffer.
+        unsafe { writer.copy_slice_unchecked(&source[start..][..length]) };
     }
-
-    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) };
+    writer.finish()?;
     Ok(views.freeze())
 }
 
@@ -202,21 +197,13 @@ where
     L: UnsignedPType,
 {
     let mut views = BufferMut::<BinaryView>::with_capacity(output_len);
-    let mut cursor = 0usize;
+    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_();
-        cursor = copy_slice_to_spare(&mut views, cursor, &source[start..][..length], output_len)?;
+        writer.copy_slice(&source[start..][..length])?;
     }
-
-    vortex_ensure!(
-        cursor == output_len,
-        "PiecewiseSequenceArray expanded length {} does not match declared length {output_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) };
+    writer.finish()?;
     Ok(views.freeze())
 }