From 6b75bbaed893c272a0fcb6248f2a460c58ae24cf Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 15:33:07 -0400 Subject: [PATCH 1/8] Use cursor copies for PiecewiseSequence slices Signed-off-by: Daniel King --- .../src/arrays/decimal/compute/take.rs | 21 ++++++++--- .../src/arrays/piecewise_sequence/mod.rs | 36 +++++++++++++++++++ .../src/arrays/primitive/compute/take/mod.rs | 21 ++++++++--- .../src/arrays/varbin/compute/take.rs | 31 ++++++++++++++-- .../src/arrays/varbinview/compute/take.rs | 21 ++++++++--- 5 files changed, 116 insertions(+), 14 deletions(-) diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index 31523c02a8a..d129122968d 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -18,6 +18,7 @@ use crate::arrays::PiecewiseSequence; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; 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; @@ -199,11 +200,19 @@ where ); let mut result = BufferMut::::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()) } @@ -219,17 +228,21 @@ where T: NativeDecimalType, { let mut result = BufferMut::::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()) } diff --git a/vortex-array/src/arrays/piecewise_sequence/mod.rs b/vortex-array/src/arrays/piecewise_sequence/mod.rs index 05b7b13edf4..a6459804997 100644 --- a/vortex-array/src/arrays/piecewise_sequence/mod.rs +++ b/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; @@ -92,6 +94,40 @@ pub(crate) fn is_constant_one(multipliers: &ArrayRef) -> bool { ) } +pub(crate) fn copy_slice_to_spare( + buffer: &mut BufferMut, + cursor: usize, + source: &[T], + output_len: usize, +) -> VortexResult { + 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) +} + pub(crate) fn constant_unsigned_usize(array: &ConstantArray) -> usize { let pvalue = array .scalar() diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index 70401fcd367..d0ea4ec7de1 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -25,6 +25,7 @@ use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; 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; @@ -245,17 +246,21 @@ where L: UnsignedPType, { let mut values = BufferMut::::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()) } @@ -312,11 +317,19 @@ where ); let mut values = BufferMut::::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()) } diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index f69ed43e8fa..8e58fa0b7e1 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -22,6 +22,7 @@ use crate::arrays::VarBin; use crate::arrays::VarBinArray; use crate::arrays::dict::TakeExecute; 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; @@ -474,6 +475,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 { @@ -483,8 +485,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) @@ -551,6 +565,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_(); @@ -561,8 +576,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) diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index 5af1c1b032c..873544df6a6 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -24,6 +24,7 @@ use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; use crate::arrays::dict::TakeExecute; 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; @@ -174,11 +175,19 @@ where ); let mut views = BufferMut::::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()) } @@ -193,17 +202,21 @@ where L: UnsignedPType, { let mut views = BufferMut::::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()) } From c8b7810b7d282b656fe93ef05b44ccbbe351a654 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Mon, 20 Jul 2026 16:02:32 -0400 Subject: [PATCH 2/8] Advance PiecewiseSequence copy destinations Signed-off-by: Daniel King --- .../src/arrays/decimal/compute/take.rs | 31 ++---- .../src/arrays/piecewise_sequence/mod.rs | 98 +++++++++++++------ .../src/arrays/primitive/compute/take/mod.rs | 31 ++---- .../src/arrays/varbin/compute/take.rs | 38 ++----- .../src/arrays/varbinview/compute/take.rs | 31 ++---- 5 files changed, 105 insertions(+), 124 deletions(-) diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index d129122968d..55d8d890e05 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/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::::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::::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()) } diff --git a/vortex-array/src/arrays/piecewise_sequence/mod.rs b/vortex-array/src/arrays/piecewise_sequence/mod.rs index a6459804997..f7d47bb54b2 100644 --- a/vortex-array/src/arrays/piecewise_sequence/mod.rs +++ b/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( - buffer: &mut BufferMut, - cursor: usize, - source: &[T], +pub(crate) struct SpareBufferWriter<'a, T> { + buffer: &'a mut BufferMut, + next: *mut T, + remaining: usize, output_len: usize, -) -> VortexResult { - 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, output_len: usize) -> VortexResult { + 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 { diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index d0ea4ec7de1..5f5b83c02da 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/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::::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::::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()) } diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 8e58fa0b7e1..89272ba74ca 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/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) diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index 873544df6a6..6a8a4f2ce94 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/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::::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::::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()) } From 9443f058d816989db3a2b05d48063e713209538d Mon Sep 17 00:00:00 2001 From: Daniel King Date: Mon, 20 Jul 2026 16:26:31 -0400 Subject: [PATCH 3/8] Move spare buffer writer to vortex-buffer Signed-off-by: Daniel King --- .../src/arrays/decimal/compute/take.rs | 2 +- .../src/arrays/piecewise_sequence/mod.rs | 76 ---------- .../src/arrays/primitive/compute/take/mod.rs | 2 +- .../src/arrays/varbin/compute/take.rs | 2 +- .../src/arrays/varbinview/compute/take.rs | 2 +- vortex-buffer/src/lib.rs | 2 + vortex-buffer/src/spare_buffer_writer.rs | 143 ++++++++++++++++++ 7 files changed, 149 insertions(+), 80 deletions(-) create mode 100644 vortex-buffer/src/spare_buffer_writer.rs diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index 55d8d890e05..14983536213 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -4,6 +4,7 @@ use itertools::Itertools as _; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; +use vortex_buffer::SpareBufferWriter; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; @@ -17,7 +18,6 @@ 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; diff --git a/vortex-array/src/arrays/piecewise_sequence/mod.rs b/vortex-array/src/arrays/piecewise_sequence/mod.rs index f7d47bb54b2..05b7b13edf4 100644 --- a/vortex-array/src/arrays/piecewise_sequence/mod.rs +++ b/vortex-array/src/arrays/piecewise_sequence/mod.rs @@ -8,8 +8,6 @@ //! 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; @@ -94,80 +92,6 @@ pub(crate) fn is_constant_one(multipliers: &ArrayRef) -> bool { ) } -pub(crate) struct SpareBufferWriter<'a, T> { - buffer: &'a mut BufferMut, - next: *mut T, - remaining: usize, - output_len: usize, -} - -impl<'a, T: Copy> SpareBufferWriter<'a, T> { - pub(crate) fn new(buffer: &'a mut BufferMut, output_len: usize) -> VortexResult { - 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() - ); - - 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(()) - } -} - pub(crate) fn constant_unsigned_usize(array: &ConstantArray) -> usize { let pvalue = array .scalar() diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index 5f5b83c02da..7675963cc24 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -9,6 +9,7 @@ use std::sync::LazyLock; use itertools::Itertools as _; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; +use vortex_buffer::SpareBufferWriter; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -24,7 +25,6 @@ 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; diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 89272ba74ca..548a4cb9cce 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -5,6 +5,7 @@ use itertools::Itertools as _; use vortex_buffer::BitBufferMut; use vortex_buffer::BufferMut; use vortex_buffer::ByteBufferMut; +use vortex_buffer::SpareBufferWriter; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -21,7 +22,6 @@ 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; diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index 6a8a4f2ce94..f0d40eda1d0 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -8,6 +8,7 @@ use itertools::Itertools as _; use num_traits::AsPrimitive; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; +use vortex_buffer::SpareBufferWriter; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; @@ -23,7 +24,6 @@ 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; diff --git a/vortex-buffer/src/lib.rs b/vortex-buffer/src/lib.rs index ee113481353..7fc58fff75b 100644 --- a/vortex-buffer/src/lib.rs +++ b/vortex-buffer/src/lib.rs @@ -53,6 +53,7 @@ pub use buffer_mut::*; pub use bytes::*; pub use r#const::*; pub use dispatch::*; +pub use spare_buffer_writer::*; pub use string::*; mod alignment; #[cfg(feature = "arrow")] @@ -69,6 +70,7 @@ mod macros; mod memmap2; #[cfg(feature = "serde")] mod serde; +mod spare_buffer_writer; mod string; /// Trusted-length iterator trait and adapters for safe pre-allocation. pub mod trusted_len; diff --git a/vortex-buffer/src/spare_buffer_writer.rs b/vortex-buffer/src/spare_buffer_writer.rs new file mode 100644 index 00000000000..9b56e69eccc --- /dev/null +++ b/vortex-buffer/src/spare_buffer_writer.rs @@ -0,0 +1,143 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ptr; + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::BufferMut; + +/// 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. +pub struct SpareBufferWriter<'a, T> { + buffer: &'a mut BufferMut, + next: *mut T, + remaining: 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 fn new(buffer: &'a mut BufferMut, output_len: usize) -> VortexResult { + 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() + ); + + let next = buffer.spare_capacity_mut().as_mut_ptr().cast(); + Ok(Self { + buffer, + next, + remaining: output_len, + output_len, + }) + } + + /// Copies `source` into the next output slots. + #[inline] + pub 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 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(); + } + + /// Sets the target buffer length after exactly `output_len` values have been written. + pub 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(()) + } +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use super::SpareBufferWriter; + use crate::BufferMut; + + #[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::::new(&mut buffer, 2)?; + let error = writer.finish().unwrap_err(); + + assert!( + error + .to_string() + .contains("does not match declared output length") + ); + Ok(()) + } +} From 4395d389776710934a8742402c03ff7754d0ae38 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 20 Jul 2026 22:53:13 +0000 Subject: [PATCH 4/8] Make SpareBufferWriter bounds-checked and drop caller-side unsafe Replace the raw-pointer cursor with a written-count cursor over spare_capacity_mut, and remove copy_slice_unchecked along with the unsafe blocks at the four PiecewiseSequence call sites. The bounds check is one well-predicted branch per slice: benchmarking the gather loop over 16M u64 elements shows checked copies within noise of unchecked ones at every run length (212ms vs 200ms even at the pathological run length of 1, where both beat extend_from_slice's 299ms), so the unchecked variant's extra unsafe surface isn't paying for anything. The writer is now #[must_use], and the remaining unsafety is contained to the copy into the bounds-checked spare-capacity range (the not-yet-stable <[MaybeUninit]>::write_copy_of_slice) and the single set_len in finish, justified by the writer's contiguous-prefix invariant. Signed-off-by: "Robert" Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_01GUSY34eRKq9du5zQK1FuZq --- .../src/arrays/decimal/compute/take.rs | 4 +- .../src/arrays/primitive/compute/take/mod.rs | 4 +- .../src/arrays/varbin/compute/take.rs | 6 +-- .../src/arrays/varbinview/compute/take.rs | 4 +- vortex-buffer/src/spare_buffer_writer.rs | 54 +++++++++---------- 5 files changed, 31 insertions(+), 41 deletions(-) diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index 14983536213..487649cdeec 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -203,9 +203,7 @@ where let mut writer = SpareBufferWriter::new(&mut result, output_len)?; for &start in starts { let start = start.as_(); - // SAFETY: `computed_len == output_len` proves that all fixed-length slices fit in the - // output buffer. - unsafe { writer.copy_slice_unchecked(&values[start..][..length]) }; + writer.copy_slice(&values[start..][..length])?; } writer.finish()?; Ok(result.freeze()) diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index 7675963cc24..bdc5e55e767 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -312,9 +312,7 @@ where let mut writer = SpareBufferWriter::new(&mut values, output_len)?; for &start in starts { let start = start.as_(); - // SAFETY: `computed_len == output_len` proves that all fixed-length slices fit in the - // output buffer. - unsafe { writer.copy_slice_unchecked(&source[start..][..length]) }; + writer.copy_slice(&source[start..][..length])?; } writer.finish()?; Ok(values.freeze()) diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 548a4cb9cce..c76b3633ac8 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -485,8 +485,7 @@ where let offset_range = &offsets[start..][..=length]; let byte_start = offset_range[0].as_(); let byte_end = offset_range[length].as_(); - // SAFETY: the first pass computed `output_bytes` from these same byte ranges. - unsafe { writer.copy_slice_unchecked(&data[byte_start..][..byte_end - byte_start]) }; + writer.copy_slice(&data[byte_start..][..byte_end - byte_start])?; } writer.finish()?; @@ -566,8 +565,7 @@ where let offset_range = &offsets[start..][..=length]; let byte_start = offset_range[0].as_(); let byte_end = offset_range[length].as_(); - // SAFETY: the first pass computed `output_bytes` from these same byte ranges. - unsafe { writer.copy_slice_unchecked(&data[byte_start..][..byte_end - byte_start]) }; + writer.copy_slice(&data[byte_start..][..byte_end - byte_start])?; } writer.finish()?; diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index f0d40eda1d0..1128dc03705 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -178,9 +178,7 @@ where let mut writer = SpareBufferWriter::new(&mut views, output_len)?; for &start in starts { let start = start.as_(); - // SAFETY: `computed_len == output_len` proves that all fixed-length slices fit in the - // output buffer. - unsafe { writer.copy_slice_unchecked(&source[start..][..length]) }; + writer.copy_slice(&source[start..][..length])?; } writer.finish()?; Ok(views.freeze()) diff --git a/vortex-buffer/src/spare_buffer_writer.rs b/vortex-buffer/src/spare_buffer_writer.rs index 9b56e69eccc..097381fa7ad 100644 --- a/vortex-buffer/src/spare_buffer_writer.rs +++ b/vortex-buffer/src/spare_buffer_writer.rs @@ -12,10 +12,19 @@ use crate::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 struct SpareBufferWriter<'a, T> { buffer: &'a mut BufferMut, - next: *mut T, - remaining: usize, + written: usize, output_len: usize, } @@ -35,11 +44,9 @@ impl<'a, T: Copy> SpareBufferWriter<'a, T> { buffer.capacity() ); - let next = buffer.spare_capacity_mut().as_mut_ptr().cast(); Ok(Self { buffer, - next, - remaining: output_len, + written: 0, output_len, }) } @@ -48,42 +55,33 @@ impl<'a, T: Copy> SpareBufferWriter<'a, T> { #[inline] pub fn copy_slice(&mut self, source: &[T]) -> VortexResult<()> { vortex_ensure!( - source.len() <= self.remaining, + source.len() <= self.output_len - self.written, "slice copy length {} exceeds remaining output length {}", source.len(), - self.remaining + self.output_len - self.written ); - // 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 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. + let end = self.written + source.len(); + let dst = &mut self.buffer.spare_capacity_mut()[self.written..end]; + // SAFETY: `MaybeUninit` 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]>::write_copy_of_slice`, which is not yet stable. unsafe { - ptr::copy_nonoverlapping(source.as_ptr(), self.next, source.len()); - self.next = self.next.add(source.len()); + ptr::copy_nonoverlapping(source.as_ptr(), dst.as_mut_ptr().cast::(), source.len()); } - self.remaining -= source.len(); + self.written = end; + Ok(()) } /// Sets the target buffer length after exactly `output_len` values have been written. pub fn finish(self) -> VortexResult<()> { vortex_ensure!( - self.remaining == 0, + self.written == self.output_len, "slice copy length {} does not match declared output length {}", - self.output_len - self.remaining, + self.written, self.output_len ); - // SAFETY: successful calls to the copy methods initialized exactly `output_len` slots. + // 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); } From 8235e0382439433549f5d8e61b3932bc997c1b16 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 20 Jul 2026 23:15:38 +0000 Subject: [PATCH 5/8] Move SpareBufferWriter back into vortex-array The writer only exists to serve the PiecewiseSequence gather loops, so keep it pub(crate) in arrays::piecewise_sequence as its own submodule rather than growing vortex-buffer's public API. It only uses BufferMut's public surface, so nothing in vortex-buffer changes behavior. Signed-off-by: "Robert" Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_01GUSY34eRKq9du5zQK1FuZq --- vortex-array/src/arrays/decimal/compute/take.rs | 2 +- vortex-array/src/arrays/piecewise_sequence/mod.rs | 2 ++ .../piecewise_sequence}/spare_buffer_writer.rs | 13 ++++++------- .../src/arrays/primitive/compute/take/mod.rs | 2 +- vortex-array/src/arrays/varbin/compute/take.rs | 2 +- vortex-array/src/arrays/varbinview/compute/take.rs | 2 +- vortex-buffer/src/lib.rs | 2 -- 7 files changed, 12 insertions(+), 13 deletions(-) rename {vortex-buffer/src => vortex-array/src/arrays/piecewise_sequence}/spare_buffer_writer.rs (93%) diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index 487649cdeec..b897afd562e 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -4,7 +4,6 @@ use itertools::Itertools as _; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; -use vortex_buffer::SpareBufferWriter; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; @@ -18,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; diff --git a/vortex-array/src/arrays/piecewise_sequence/mod.rs b/vortex-array/src/arrays/piecewise_sequence/mod.rs index 05b7b13edf4..98c0e6d3501 100644 --- a/vortex-array/src/arrays/piecewise_sequence/mod.rs +++ b/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( diff --git a/vortex-buffer/src/spare_buffer_writer.rs b/vortex-array/src/arrays/piecewise_sequence/spare_buffer_writer.rs similarity index 93% rename from vortex-buffer/src/spare_buffer_writer.rs rename to vortex-array/src/arrays/piecewise_sequence/spare_buffer_writer.rs index 097381fa7ad..2cacf09000b 100644 --- a/vortex-buffer/src/spare_buffer_writer.rs +++ b/vortex-array/src/arrays/piecewise_sequence/spare_buffer_writer.rs @@ -3,11 +3,10 @@ use std::ptr; +use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use crate::BufferMut; - /// 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 @@ -22,7 +21,7 @@ use crate::BufferMut; /// [`set_len`]: BufferMut::set_len /// [`finish`]: SpareBufferWriter::finish #[must_use = "call `finish` to set the buffer length after writing"] -pub struct SpareBufferWriter<'a, T> { +pub(crate) struct SpareBufferWriter<'a, T> { buffer: &'a mut BufferMut, written: usize, output_len: usize, @@ -32,7 +31,7 @@ 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 fn new(buffer: &'a mut BufferMut, output_len: usize) -> VortexResult { + pub(crate) fn new(buffer: &'a mut BufferMut, output_len: usize) -> VortexResult { vortex_ensure!( buffer.is_empty(), "slice copy buffer already has {} initialized values", @@ -53,7 +52,7 @@ impl<'a, T: Copy> SpareBufferWriter<'a, T> { /// Copies `source` into the next output slots. #[inline] - pub fn copy_slice(&mut self, source: &[T]) -> VortexResult<()> { + 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 {}", @@ -73,7 +72,7 @@ impl<'a, T: Copy> SpareBufferWriter<'a, T> { } /// Sets the target buffer length after exactly `output_len` values have been written. - pub fn finish(self) -> VortexResult<()> { + pub(crate) fn finish(self) -> VortexResult<()> { vortex_ensure!( self.written == self.output_len, "slice copy length {} does not match declared output length {}", @@ -91,10 +90,10 @@ impl<'a, T: Copy> SpareBufferWriter<'a, T> { #[cfg(test)] mod tests { + use vortex_buffer::BufferMut; use vortex_error::VortexResult; use super::SpareBufferWriter; - use crate::BufferMut; #[test] fn writes_slices_into_spare_capacity() -> VortexResult<()> { diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index bdc5e55e767..442279e1dfb 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -9,7 +9,6 @@ use std::sync::LazyLock; use itertools::Itertools as _; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; -use vortex_buffer::SpareBufferWriter; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -25,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; diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index c76b3633ac8..9723c9a4bba 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -5,7 +5,6 @@ use itertools::Itertools as _; use vortex_buffer::BitBufferMut; use vortex_buffer::BufferMut; use vortex_buffer::ByteBufferMut; -use vortex_buffer::SpareBufferWriter; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -22,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; diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index 1128dc03705..b7559a32fb6 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -8,7 +8,6 @@ use itertools::Itertools as _; use num_traits::AsPrimitive; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; -use vortex_buffer::SpareBufferWriter; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; @@ -24,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; diff --git a/vortex-buffer/src/lib.rs b/vortex-buffer/src/lib.rs index 7fc58fff75b..ee113481353 100644 --- a/vortex-buffer/src/lib.rs +++ b/vortex-buffer/src/lib.rs @@ -53,7 +53,6 @@ pub use buffer_mut::*; pub use bytes::*; pub use r#const::*; pub use dispatch::*; -pub use spare_buffer_writer::*; pub use string::*; mod alignment; #[cfg(feature = "arrow")] @@ -70,7 +69,6 @@ mod macros; mod memmap2; #[cfg(feature = "serde")] mod serde; -mod spare_buffer_writer; mod string; /// Trusted-length iterator trait and adapters for safe pre-allocation. pub mod trusted_len; From 05ce4081a6b9525b3f80cae781d0439f48f19620 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 21 Jul 2026 14:06:24 +0100 Subject: [PATCH 6/8] less Signed-off-by: Robert Kruszewski --- vortex-array/src/arrays/bool/compute/take.rs | 2 +- .../src/arrays/decimal/compute/take.rs | 40 ++++- vortex-array/src/arrays/list/compute/take.rs | 8 +- .../src/arrays/piecewise_sequence/mod.rs | 2 - .../piecewise_sequence/spare_buffer_writer.rs | 140 ------------------ .../src/arrays/primitive/compute/take/mod.rs | 44 +++++- .../src/arrays/varbin/compute/take.rs | 45 +++++- .../src/arrays/varbinview/compute/take.rs | 50 ++++++- 8 files changed, 155 insertions(+), 176 deletions(-) delete mode 100644 vortex-array/src/arrays/piecewise_sequence/spare_buffer_writer.rs diff --git a/vortex-array/src/arrays/bool/compute/take.rs b/vortex-array/src/arrays/bool/compute/take.rs index 14d90a5e7b1..529450f6639 100644 --- a/vortex-array/src/arrays/bool/compute/take.rs +++ b/vortex-array/src/arrays/bool/compute/take.rs @@ -159,7 +159,7 @@ where ); let mut values = BitBufferMut::with_capacity(output_len); - for &start in starts { + for start in starts { let start = start.as_(); values.append_buffer(&source.slice(start..).slice(..length)); } diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index b897afd562e..170309ee32c 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ptr; + use itertools::Itertools as _; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -17,7 +19,6 @@ 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; @@ -200,12 +201,25 @@ where ); let mut result = BufferMut::::with_capacity(output_len); - let mut writer = SpareBufferWriter::new(&mut result, output_len)?; + let spare = &mut result.spare_capacity_mut()[..output_len]; + let mut cursor = 0usize; for &start in starts { let start = start.as_(); - writer.copy_slice(&values[start..][..length])?; + let src = &values[start..][..length]; + // SAFETY: `MaybeUninit` has the same layout as `T`, source and destination have equal + // lengths, and `spare` exclusively borrows the output buffer so they cannot overlap. + unsafe { + ptr::copy_nonoverlapping( + src.as_ptr(), + spare[cursor..][..src.len()].as_mut_ptr().cast::(), + src.len(), + ); + } + cursor += src.len(); } - writer.finish()?; + // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity, and + // `computed_len == output_len` proves the loop filled exactly `output_len` slots. + unsafe { result.set_len(cursor) }; Ok(result.freeze()) } @@ -221,13 +235,25 @@ where T: NativeDecimalType, { let mut result = BufferMut::::with_capacity(output_len); - let mut writer = SpareBufferWriter::new(&mut result, output_len)?; + let spare = &mut result.spare_capacity_mut()[..output_len]; + let mut cursor = 0usize; for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); - writer.copy_slice(&values[start..][..length])?; + let src = &values[start..][..length]; + // SAFETY: `MaybeUninit` has the same layout as `T`, source and destination have equal + // lengths, and `spare` exclusively borrows the output buffer so they cannot overlap. + unsafe { + ptr::copy_nonoverlapping( + src.as_ptr(), + spare[cursor..][..src.len()].as_mut_ptr().cast::(), + src.len(), + ); + } + cursor += src.len(); } - writer.finish()?; + // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity. + unsafe { result.set_len(cursor) }; Ok(result.freeze()) } diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index 8cb5c8707c3..7ac16f0dd1c 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -463,7 +463,7 @@ where } let mut total = 0usize; - for &start in starts { + for start in starts { let start: usize = start.as_(); let offset_range = &offsets[start..][..=length]; let element_start: usize = offset_range[0].as_(); @@ -490,7 +490,7 @@ where } let mut total = 0usize; - for &start in starts { + for start in starts { let start: usize = start.as_(); let additional = valid_piece_elements_len(offsets, data_validity, start, length)?; total = total @@ -593,7 +593,7 @@ where let mut output_elements = 0usize; new_offsets.push(OutputOffset::zero()); - for &start in starts { + for start in starts { let start: usize = start.as_(); if length == 0 { continue; @@ -658,7 +658,7 @@ where }; gather.new_offsets.push(OutputOffset::zero()); - for &start in starts { + for start in starts { let start: usize = start.as_(); if length == 0 { continue; diff --git a/vortex-array/src/arrays/piecewise_sequence/mod.rs b/vortex-array/src/arrays/piecewise_sequence/mod.rs index 98c0e6d3501..05b7b13edf4 100644 --- a/vortex-array/src/arrays/piecewise_sequence/mod.rs +++ b/vortex-array/src/arrays/piecewise_sequence/mod.rs @@ -28,14 +28,12 @@ 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( diff --git a/vortex-array/src/arrays/piecewise_sequence/spare_buffer_writer.rs b/vortex-array/src/arrays/piecewise_sequence/spare_buffer_writer.rs deleted file mode 100644 index 2cacf09000b..00000000000 --- a/vortex-array/src/arrays/piecewise_sequence/spare_buffer_writer.rs +++ /dev/null @@ -1,140 +0,0 @@ -// 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, - 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, output_len: usize) -> VortexResult { - 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` 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]>::write_copy_of_slice`, which is not yet stable. - unsafe { - ptr::copy_nonoverlapping(source.as_ptr(), dst.as_mut_ptr().cast::(), 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::::new(&mut buffer, 2)?; - let error = writer.finish().unwrap_err(); - - assert!( - error - .to_string() - .contains("does not match declared output length") - ); - Ok(()) - } -} diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index 442279e1dfb..4e4699d7090 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -4,6 +4,7 @@ #[cfg(any(target_arch = "x86_64", target_arch = "x86"))] mod avx2; +use std::ptr; use std::sync::LazyLock; use itertools::Itertools as _; @@ -24,7 +25,6 @@ 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; @@ -246,13 +246,30 @@ where L: UnsignedPType, { let mut values = BufferMut::::with_capacity(output_len); - let mut writer = SpareBufferWriter::new(&mut values, output_len)?; + let spare = &mut values.spare_capacity_mut()[..output_len]; + let mut cursor = 0usize; for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); - writer.copy_slice(&source[start..][..length])?; + let src = &source[start..][..length]; + // SAFETY: `MaybeUninit` has the same layout as `T`, source and destination have equal + // lengths, and `spare` exclusively borrows the output buffer so they cannot overlap. + unsafe { + ptr::copy_nonoverlapping( + src.as_ptr(), + spare[cursor..][..src.len()].as_mut_ptr().cast::(), + src.len(), + ); + } + cursor += src.len(); } - writer.finish()?; + // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity. + unsafe { values.set_len(cursor) }; + vortex_ensure!( + values.len() == output_len, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + values.len() + ); Ok(values.freeze()) } @@ -309,12 +326,25 @@ where ); let mut values = BufferMut::::with_capacity(output_len); - let mut writer = SpareBufferWriter::new(&mut values, output_len)?; + let spare = &mut values.spare_capacity_mut()[..output_len]; + let mut cursor = 0usize; for &start in starts { let start = start.as_(); - writer.copy_slice(&source[start..][..length])?; + let src = &source[start..][..length]; + // SAFETY: `MaybeUninit` has the same layout as `T`, source and destination have equal + // lengths, and `spare` exclusively borrows the output buffer so they cannot overlap. + unsafe { + ptr::copy_nonoverlapping( + src.as_ptr(), + spare[cursor..][..src.len()].as_mut_ptr().cast::(), + src.len(), + ); + } + cursor += src.len(); } - writer.finish()?; + // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity, and + // `computed_len == output_len` proves the loop filled exactly `output_len` slots. + unsafe { values.set_len(cursor) }; Ok(values.freeze()) } diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 9723c9a4bba..67bf05e858e 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ptr; + use itertools::Itertools as _; use vortex_buffer::BitBufferMut; use vortex_buffer::BufferMut; @@ -21,7 +23,6 @@ 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; @@ -443,7 +444,7 @@ where new_offsets.push(NewOffset::zero()); let mut output_bytes = 0usize; - for &start in starts { + for start in starts { let start = start.as_(); if length == 0 { continue; @@ -475,7 +476,8 @@ where } let mut new_data = ByteBufferMut::with_capacity(output_bytes); - let mut writer = SpareBufferWriter::new(&mut new_data, output_bytes)?; + let spare = &mut new_data.spare_capacity_mut()[..output_bytes]; + let mut cursor = 0usize; for &start in starts { let start = start.as_(); if length == 0 { @@ -485,9 +487,22 @@ where let offset_range = &offsets[start..][..=length]; let byte_start = offset_range[0].as_(); let byte_end = offset_range[length].as_(); - writer.copy_slice(&data[byte_start..][..byte_end - byte_start])?; + let src = &data[byte_start..][..byte_end - byte_start]; + // SAFETY: `MaybeUninit` has the same layout as `u8`, source and destination have + // equal lengths, and `spare` exclusively borrows the output buffer so they cannot + // overlap. + unsafe { + ptr::copy_nonoverlapping( + src.as_ptr(), + spare[cursor..][..src.len()].as_mut_ptr().cast::(), + src.len(), + ); + } + cursor += src.len(); } - writer.finish()?; + // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity, and the first + // pass over the same offsets computed `output_bytes` as the sum of the copied slice lengths. + unsafe { new_data.set_len(cursor) }; let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) .reinterpret_cast(out_offset_ptype) @@ -554,7 +569,8 @@ where ); let mut new_data = ByteBufferMut::with_capacity(output_bytes); - let mut writer = SpareBufferWriter::new(&mut new_data, output_bytes)?; + let spare = &mut new_data.spare_capacity_mut()[..output_bytes]; + let mut cursor = 0usize; for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); @@ -565,9 +581,22 @@ where let offset_range = &offsets[start..][..=length]; let byte_start = offset_range[0].as_(); let byte_end = offset_range[length].as_(); - writer.copy_slice(&data[byte_start..][..byte_end - byte_start])?; + let src = &data[byte_start..byte_end]; + // SAFETY: `MaybeUninit` has the same layout as `u8`, source and destination have + // equal lengths, and `spare` exclusively borrows the output buffer so they cannot + // overlap. + unsafe { + ptr::copy_nonoverlapping( + src.as_ptr(), + spare[cursor..][..src.len()].as_mut_ptr().cast::(), + src.len(), + ); + } + cursor += src.len(); } - writer.finish()?; + // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity, and the first + // pass over the same offsets computed `output_bytes` as the sum of the copied slice lengths. + unsafe { new_data.set_len(cursor) }; let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) .reinterpret_cast(out_offset_ptype) diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index b7559a32fb6..46e4d30cb1b 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::iter; +use std::ptr; use std::sync::Arc; use itertools::Itertools as _; @@ -23,7 +24,6 @@ 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; @@ -175,12 +175,28 @@ where ); let mut views = BufferMut::::with_capacity(output_len); - let mut writer = SpareBufferWriter::new(&mut views, output_len)?; + let spare = &mut views.spare_capacity_mut()[..output_len]; + let mut cursor = 0usize; for &start in starts { let start = start.as_(); - writer.copy_slice(&source[start..][..length])?; + let src = &source[start..][..length]; + // SAFETY: `MaybeUninit` has the same layout as `BinaryView`, source and + // destination have equal lengths, and `spare` exclusively borrows the output buffer so + // they cannot overlap. + unsafe { + ptr::copy_nonoverlapping( + src.as_ptr(), + spare[cursor..][..src.len()] + .as_mut_ptr() + .cast::(), + src.len(), + ); + } + cursor += src.len(); } - writer.finish()?; + // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity, and + // `computed_len == output_len` proves the loop filled exactly `output_len` slots. + unsafe { views.set_len(cursor) }; Ok(views.freeze()) } @@ -195,13 +211,33 @@ where L: UnsignedPType, { let mut views = BufferMut::::with_capacity(output_len); - let mut writer = SpareBufferWriter::new(&mut views, output_len)?; + let spare = &mut views.spare_capacity_mut()[..output_len]; + let mut cursor = 0usize; for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); - writer.copy_slice(&source[start..][..length])?; + let src = &source[start..][..length]; + // SAFETY: `MaybeUninit` has the same layout as `BinaryView`, source and + // destination have equal lengths, and `spare` exclusively borrows the output buffer so + // they cannot overlap. + unsafe { + ptr::copy_nonoverlapping( + src.as_ptr(), + spare[cursor..][..src.len()] + .as_mut_ptr() + .cast::(), + src.len(), + ); + } + cursor += src.len(); } - writer.finish()?; + // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity. + unsafe { views.set_len(cursor) }; + vortex_ensure!( + views.len() == output_len, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + views.len() + ); Ok(views.freeze()) } From dffff6887c4d98ec41c9e60e83e1305ddb139d23 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 21 Jul 2026 14:12:29 +0100 Subject: [PATCH 7/8] comments Signed-off-by: Robert Kruszewski --- .../src/arrays/decimal/compute/take.rs | 19 ++++++++++----- .../src/arrays/primitive/compute/take/mod.rs | 14 ++++++----- .../src/arrays/varbin/compute/take.rs | 24 +++++++++++-------- .../src/arrays/varbinview/compute/take.rs | 16 ++++++------- 4 files changed, 43 insertions(+), 30 deletions(-) diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index 170309ee32c..d7bde288261 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -206,8 +206,7 @@ where for &start in starts { let start = start.as_(); let src = &values[start..][..length]; - // SAFETY: `MaybeUninit` has the same layout as `T`, source and destination have equal - // lengths, and `spare` exclusively borrows the output buffer so they cannot overlap. + // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. unsafe { ptr::copy_nonoverlapping( src.as_ptr(), @@ -217,9 +216,13 @@ where } cursor += src.len(); } - // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity, and - // `computed_len == output_len` proves the loop filled exactly `output_len` slots. + // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity. unsafe { result.set_len(cursor) }; + vortex_ensure!( + result.len() == output_len, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + result.len() + ); Ok(result.freeze()) } @@ -241,8 +244,7 @@ where let start = start.as_(); let length = length.as_(); let src = &values[start..][..length]; - // SAFETY: `MaybeUninit` has the same layout as `T`, source and destination have equal - // lengths, and `spare` exclusively borrows the output buffer so they cannot overlap. + // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. unsafe { ptr::copy_nonoverlapping( src.as_ptr(), @@ -254,6 +256,11 @@ where } // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity. unsafe { result.set_len(cursor) }; + vortex_ensure!( + result.len() == output_len, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + result.len() + ); Ok(result.freeze()) } diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index 4e4699d7090..5d663f55967 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -252,8 +252,7 @@ where let start = start.as_(); let length = length.as_(); let src = &source[start..][..length]; - // SAFETY: `MaybeUninit` has the same layout as `T`, source and destination have equal - // lengths, and `spare` exclusively borrows the output buffer so they cannot overlap. + // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. unsafe { ptr::copy_nonoverlapping( src.as_ptr(), @@ -331,8 +330,7 @@ where for &start in starts { let start = start.as_(); let src = &source[start..][..length]; - // SAFETY: `MaybeUninit` has the same layout as `T`, source and destination have equal - // lengths, and `spare` exclusively borrows the output buffer so they cannot overlap. + // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. unsafe { ptr::copy_nonoverlapping( src.as_ptr(), @@ -342,9 +340,13 @@ where } cursor += src.len(); } - // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity, and - // `computed_len == output_len` proves the loop filled exactly `output_len` slots. + // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity. unsafe { values.set_len(cursor) }; + vortex_ensure!( + values.len() == output_len, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + values.len() + ); Ok(values.freeze()) } diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 67bf05e858e..2a72dad26be 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -488,9 +488,7 @@ where let byte_start = offset_range[0].as_(); let byte_end = offset_range[length].as_(); let src = &data[byte_start..][..byte_end - byte_start]; - // SAFETY: `MaybeUninit` has the same layout as `u8`, source and destination have - // equal lengths, and `spare` exclusively borrows the output buffer so they cannot - // overlap. + // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. unsafe { ptr::copy_nonoverlapping( src.as_ptr(), @@ -500,9 +498,13 @@ where } cursor += src.len(); } - // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity, and the first - // pass over the same offsets computed `output_bytes` as the sum of the copied slice lengths. + // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity. unsafe { new_data.set_len(cursor) }; + vortex_ensure!( + new_data.len() == output_bytes, + "PiecewiseSequenceArray gathered byte length {} does not match declared byte length {output_bytes}", + new_data.len() + ); let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) .reinterpret_cast(out_offset_ptype) @@ -582,9 +584,7 @@ where let byte_start = offset_range[0].as_(); let byte_end = offset_range[length].as_(); let src = &data[byte_start..byte_end]; - // SAFETY: `MaybeUninit` has the same layout as `u8`, source and destination have - // equal lengths, and `spare` exclusively borrows the output buffer so they cannot - // overlap. + // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. unsafe { ptr::copy_nonoverlapping( src.as_ptr(), @@ -594,9 +594,13 @@ where } cursor += src.len(); } - // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity, and the first - // pass over the same offsets computed `output_bytes` as the sum of the copied slice lengths. + // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity. unsafe { new_data.set_len(cursor) }; + vortex_ensure!( + new_data.len() == output_bytes, + "PiecewiseSequenceArray gathered byte length {} does not match declared byte length {output_bytes}", + new_data.len() + ); let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) .reinterpret_cast(out_offset_ptype) diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index 46e4d30cb1b..c10b16cd419 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -180,9 +180,7 @@ where for &start in starts { let start = start.as_(); let src = &source[start..][..length]; - // SAFETY: `MaybeUninit` has the same layout as `BinaryView`, source and - // destination have equal lengths, and `spare` exclusively borrows the output buffer so - // they cannot overlap. + // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. unsafe { ptr::copy_nonoverlapping( src.as_ptr(), @@ -194,9 +192,13 @@ where } cursor += src.len(); } - // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity, and - // `computed_len == output_len` proves the loop filled exactly `output_len` slots. + // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity. unsafe { views.set_len(cursor) }; + vortex_ensure!( + views.len() == output_len, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + views.len() + ); Ok(views.freeze()) } @@ -217,9 +219,7 @@ where let start = start.as_(); let length = length.as_(); let src = &source[start..][..length]; - // SAFETY: `MaybeUninit` has the same layout as `BinaryView`, source and - // destination have equal lengths, and `spare` exclusively borrows the output buffer so - // they cannot overlap. + // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. unsafe { ptr::copy_nonoverlapping( src.as_ptr(), From 542e9024cde9885d6519490f14f80980570970ef Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 21 Jul 2026 14:14:07 +0100 Subject: [PATCH 8/8] nit Signed-off-by: Robert Kruszewski --- vortex-array/src/arrays/varbin/compute/take.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 2a72dad26be..bc4b36cd221 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -478,7 +478,7 @@ where let mut new_data = ByteBufferMut::with_capacity(output_bytes); let spare = &mut new_data.spare_capacity_mut()[..output_bytes]; let mut cursor = 0usize; - for &start in starts { + for start in starts { let start = start.as_(); if length == 0 { continue; @@ -487,7 +487,7 @@ where let offset_range = &offsets[start..][..=length]; let byte_start = offset_range[0].as_(); let byte_end = offset_range[length].as_(); - let src = &data[byte_start..][..byte_end - byte_start]; + let src = &data[byte_start..byte_end]; // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. unsafe { ptr::copy_nonoverlapping(