Rich-T-kid commented on code in PR #11174:
URL: https://github.com/apache/arrow-rs/pull/11174#discussion_r4552818434
##########
arrow-select/src/coalesce/byte_view.rs:
##########
@@ -58,6 +58,12 @@ pub(crate) struct InProgressByteViewArray<B: ByteViewType> {
completed_buffers_size: usize,
/// The size in bytes from [`Self::source`] that it is being used in
[`Self::completed`]
size_of_completed_buffers_from_current_source: usize,
+ /// Weak identity of the last indexed source. Its data buffers are already
+ /// retained in `completed`, so keeping the array alive would retain its
+ /// views and null bitmap unnecessarily.
+ last_indexed_source: Option<Weak<dyn Array>>,
Review Comment:
Yes. It lets repeated takes from the same array reuse the buffer references
already stored in `completed`, rather than adding them again. The weak
reference avoids keeping the source views and null bitmap alive.
##########
arrow-select/src/coalesce/byte_view.rs:
##########
@@ -393,13 +401,174 @@ impl<B: ByteViewType> InProgressByteViewArray<B> {
self.views.extend(new_views);
self.current = Some(dst_buffer);
}
+
+ /// # Safety
+ ///
+ /// Every value in `indices` must be less than `src_views.len()`, and every
+ /// selected non-inline view must reference a valid range in `src_buffers`.
+ #[inline(never)]
+ unsafe fn append_indexed_views_and_copy_strings(
+ &mut self,
+ src_views: &[u128],
+ indices: &[u32],
+ src_buffers: &[Buffer],
+ selected_bytes: usize,
+ ) {
+ // SAFETY: the caller guarantees the same index and ByteView buffer
+ // invariants required by the delegated operation.
+ unsafe {
+ self.append_indexed_views_and_copy_strings_with_max_buffer_size(
+ src_views,
+ indices,
+ src_buffers,
+ selected_bytes,
+ i32::MAX as usize,
+ );
+ }
+ }
+
+ /// # Safety
+ ///
+ /// Every value in `indices` must be less than `src_views.len()`, and every
+ /// selected non-inline view must reference a valid range in `src_buffers`.
+ unsafe 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,
+ ) {
+ debug_assert!(
+ indices.iter().all(|&idx| (idx as usize) < src_views.len()),
+ "append_indexed_views_and_copy_strings_with_max_buffer_size
requires caller-validated in-bounds indices",
+ );
+ 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");
Review Comment:
Opened #12210 to track this as a follow-up.
--
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]
