Rich-T-kid commented on code in PR #11174:
URL: https://github.com/apache/arrow-rs/pull/11174#discussion_r4321707775
##########
arrow-select/src/coalesce/byte_view.rs:
##########
@@ -354,13 +362,157 @@ impl<B: ByteViewType> InProgressByteViewArray<B> {
self.views.extend(new_views);
self.current = Some(dst_buffer);
}
+
+ #[inline(never)]
+ fn append_indexed_views_and_copy_strings(
+ &mut self,
+ src_views: &[u128],
+ indices: &[u32],
+ src_buffers: &[Buffer],
+ selected_bytes: usize,
+ ) {
+ self.append_indexed_views_and_copy_strings_with_max_buffer_size(
+ src_views,
+ indices,
+ src_buffers,
+ selected_bytes,
+ i32::MAX as usize,
+ );
+ }
+
+ fn append_indexed_views_and_copy_strings_with_max_buffer_size(
+ &mut self,
+ src_views: &[u128],
+ indices: &[u32],
+ src_buffers: &[Buffer],
+ selected_bytes: usize,
+ max_buffer_size: usize,
+ ) {
+ if selected_bytes == 0 {
+ self.views.reserve(indices.len());
+ for &idx in indices {
+ // SAFETY: caller validated idx < src_views.len().
+ let raw = unsafe { *src_views.get_unchecked(idx as usize) };
+ self.views.push(raw);
+ }
+ return;
+ }
+
+ let mut dst_buffer = match self.current.take() {
+ Some(current)
+ if current.capacity() - current.len() >= selected_bytes
+ && current
+ .len()
+ .checked_add(selected_bytes)
+ .is_some_and(|len| len <= max_buffer_size) =>
+ {
+ current
+ }
+ Some(current) => {
+ let completed_buffer: Buffer = current.into();
+ self.completed_buffers_size += completed_buffer.capacity();
+ self.completed.push(completed_buffer);
+ self.buffer_source
+ .next_buffer(selected_bytes.min(max_buffer_size))
+ }
+ None => self
+ .buffer_source
+ .next_buffer(selected_bytes.min(max_buffer_size)),
+ };
+
+ let current_len = self.views.len();
+ self.views.reserve(indices.len());
+ let output =
self.views.spare_capacity_mut().as_mut_ptr().cast::<u128>();
+ let mut copied_bytes = 0usize;
+ for (i, &idx) in indices.iter().enumerate() {
+ // SAFETY: caller validated idx < src_views.len().
+ let raw = unsafe { *src_views.get_unchecked(idx as usize) };
+ let mut view = ByteView::from(raw);
+ if view.length > MAX_INLINE_VIEW_LEN {
+ let buffer_index = view.buffer_index as usize;
+ let buffer_offset = view.offset as usize;
+ let str_len = view.length as usize;
+ assert!(str_len <= max_buffer_size);
+
+ let dst_end = dst_buffer
+ .len()
+ .checked_add(str_len)
+ .expect("byte view buffer size overflow");
+ if dst_end > max_buffer_size {
+ let completed_buffer: Buffer = dst_buffer.into();
+ self.completed_buffers_size += completed_buffer.capacity();
+ self.completed.push(completed_buffer);
+
+ let remaining = selected_bytes
+ .checked_sub(copied_bytes)
+ .expect("selected byte count mismatch");
+ dst_buffer = self
+ .buffer_source
+ .next_buffer(remaining.min(max_buffer_size));
+ }
+
+ // SAFETY: source views carry valid buffer_index/offset/length.
+ let src = unsafe {
+ src_buffers
+ .get_unchecked(buffer_index)
+ .get_unchecked(buffer_offset..buffer_offset + str_len)
+ };
+ view.offset = dst_buffer
+ .len()
+ .try_into()
+ .expect("byte view offset exceeds u32::MAX");
+ view.buffer_index =
self.completed.len().try_into().expect("too many buffers");
+ let dst_len = dst_buffer.len();
+ // SAFETY: the destination was sized for the remaining
selected bytes,
+ // and `src` belongs to the input array and cannot overlap it.
+ unsafe {
+ std::ptr::copy_nonoverlapping(
+ src.as_ptr(),
+ dst_buffer.as_mut_ptr().add(dst_len),
+ str_len,
+ );
+ dst_buffer.set_len(dst_len + str_len);
+ }
Review Comment:
will swap to use `extend_from_slice` perf should be pretty much the same and
we avoid an unsafe block 👍
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
