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 31523c02a8a..d7bde288261 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; @@ -199,11 +201,28 @@ where ); let mut result = BufferMut::::with_capacity(output_len); + let spare = &mut result.spare_capacity_mut()[..output_len]; + let mut cursor = 0usize; for &start in starts { let start = start.as_(); - result.extend_from_slice(&values[start..][..length]); + let src = &values[start..][..length]; + // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. + unsafe { + ptr::copy_nonoverlapping( + src.as_ptr(), + spare[cursor..][..src.len()].as_mut_ptr().cast::(), + src.len(), + ); + } + cursor += src.len(); } - + // 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()) } @@ -219,12 +238,24 @@ where T: NativeDecimalType, { let mut result = BufferMut::::with_capacity(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_(); - result.extend_from_slice(&values[start..][..length]); + let src = &values[start..][..length]; + // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. + unsafe { + ptr::copy_nonoverlapping( + src.as_ptr(), + spare[cursor..][..src.len()].as_mut_ptr().cast::(), + src.len(), + ); + } + cursor += src.len(); } - + // 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}", 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/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index 70401fcd367..5d663f55967 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 _; @@ -245,12 +246,24 @@ where L: UnsignedPType, { let mut values = BufferMut::::with_capacity(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_(); - values.extend_from_slice(&source[start..][..length]); + let src = &source[start..][..length]; + // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. + unsafe { + ptr::copy_nonoverlapping( + src.as_ptr(), + spare[cursor..][..src.len()].as_mut_ptr().cast::(), + src.len(), + ); + } + cursor += src.len(); } - + // 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}", @@ -312,11 +325,28 @@ where ); let mut values = BufferMut::::with_capacity(output_len); + let spare = &mut values.spare_capacity_mut()[..output_len]; + let mut cursor = 0usize; for &start in starts { let start = start.as_(); - values.extend_from_slice(&source[start..][..length]); + let src = &source[start..][..length]; + // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. + unsafe { + ptr::copy_nonoverlapping( + src.as_ptr(), + spare[cursor..][..src.len()].as_mut_ptr().cast::(), + src.len(), + ); + } + cursor += src.len(); } - + // 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 f69ed43e8fa..bc4b36cd221 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; @@ -442,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; @@ -474,7 +476,9 @@ where } let mut new_data = ByteBufferMut::with_capacity(output_bytes); - for &start in starts { + 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 { continue; @@ -483,8 +487,24 @@ 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]); + let src = &data[byte_start..byte_end]; + // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. + unsafe { + ptr::copy_nonoverlapping( + src.as_ptr(), + spare[cursor..][..src.len()].as_mut_ptr().cast::(), + src.len(), + ); + } + cursor += src.len(); } + // 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) @@ -551,6 +571,8 @@ 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, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); @@ -561,8 +583,24 @@ 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]); + let src = &data[byte_start..byte_end]; + // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. + unsafe { + ptr::copy_nonoverlapping( + src.as_ptr(), + spare[cursor..][..src.len()].as_mut_ptr().cast::(), + src.len(), + ); + } + cursor += src.len(); } + // 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 5af1c1b032c..c10b16cd419 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 _; @@ -174,11 +175,30 @@ where ); let mut views = BufferMut::::with_capacity(output_len); + let spare = &mut views.spare_capacity_mut()[..output_len]; + let mut cursor = 0usize; for &start in starts { let start = start.as_(); - views.extend_from_slice(&source[start..][..length]); + let src = &source[start..][..length]; + // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. + unsafe { + ptr::copy_nonoverlapping( + src.as_ptr(), + spare[cursor..][..src.len()] + .as_mut_ptr() + .cast::(), + src.len(), + ); + } + cursor += src.len(); } - + // 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()) } @@ -193,12 +213,26 @@ where L: UnsignedPType, { let mut views = BufferMut::::with_capacity(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_(); - views.extend_from_slice(&source[start..][..length]); + let src = &source[start..][..length]; + // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. + unsafe { + ptr::copy_nonoverlapping( + src.as_ptr(), + spare[cursor..][..src.len()] + .as_mut_ptr() + .cast::(), + src.len(), + ); + } + cursor += src.len(); } - + // 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}",