Use cursor copies for PiecewiseSequence slices

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));
     }

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::<T>::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::<T>(),
+                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::<T>::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::<T>(),
+                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}",

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;

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::<T>::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::<T>(),
+                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::<T>::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::<T>(),
+                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())
 }
 

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::<u8>(),
+                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::<u8>(),
+                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)

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::<BinaryView>::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::<BinaryView>(),
+                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::<BinaryView>::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::<BinaryView>(),
+                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}",