alamb commented on code in PR #11174:
URL: https://github.com/apache/arrow-rs/pull/11174#discussion_r4551868060
##########
arrow-select/src/coalesce.rs:
##########
@@ -47,6 +50,54 @@ use generic::GenericInProgressArray;
use list::InProgressListArray;
use primitive::InProgressPrimitiveArray;
+/// Appends the validity bits selected by `indices` to `nulls`.
+///
+/// # Safety
+///
+/// `indices` must not contain nulls, and every value must be less than
+/// `source_nulls.len()`.
+pub(crate) unsafe fn append_gathered_nulls(
+ nulls: &mut NullBufferBuilder,
+ source_nulls: &NullBuffer,
+ indices: &UInt32Array,
+) {
+ debug_assert!(
+ indices
+ .values()
+ .iter()
+ .all(|&idx| (idx as usize) < source_nulls.len()),
+ "append_gathered_nulls received out-of-bounds index",
+ );
+ if source_nulls.null_count() == 0 {
+ nulls.append_n_non_nulls(indices.len());
+ return;
+ }
+ let gathered = take_bits::<UInt32Type, false>(source_nulls.inner(),
indices);
+ nulls.append_buffer(&NullBuffer::new(gathered));
+}
+
+pub(crate) fn validate_indices_bounds(
Review Comment:
This looks very similar to check_bounds
https://github.com/apache/arrow-rs/blob/b2b7d4705191a200ac7e395596c0946786d6f69c/arrow-select/src/take.rs#L180-L222
I wonder if we could reuse it instead
##########
arrow-select/src/coalesce/primitive.rs:
##########
@@ -229,6 +229,52 @@ impl<T: ArrowPrimitiveType + Debug> InProgressArray for
InProgressPrimitiveArray
}
}
+ unsafe fn copy_rows_by_indices(
+ &mut self,
+ source: &ArrayRef,
+ indices: &UInt32Array,
+ ) -> Result<(), ArrowError> {
+ if indices.is_empty() {
+ return Ok(());
+ }
+ let typed = source.as_primitive::<T>();
+ self.ensure_capacity();
+
+ if let Some(src_nulls) = typed.nulls() {
+ // SAFETY: the caller validated every index against `typed.len()`,
which
+ // equals the length of its null buffer.
+ unsafe { append_gathered_nulls(&mut self.nulls, src_nulls,
indices) };
+ } else {
+ self.nulls.append_n_non_nulls(indices.len());
+ }
+
+ let index_values = indices.values();
+ let values = typed.values();
+ let current_len = self.current.len();
+ self.current.reserve(indices.len());
+ debug_assert!(
+ index_values
+ .iter()
+ .all(|&idx| (idx as usize) < values.len()),
+ "copy_rows_by_indices requires caller-validated in-bounds indices",
+ );
+ // SAFETY: indices were bounds-checked by the coalescer, and reserve
Review Comment:
I think the safety argument now is that the caller guarnateed that the input
indices are in bounds
##########
arrow-select/src/coalesce/byte_view.rs:
##########
@@ -553,6 +722,137 @@ impl<B: ByteViewType> InProgressArray for
InProgressByteViewArray<B> {
Ok(())
}
+ unsafe fn copy_rows_by_indices(
Review Comment:
Also here it would be nice to note what the safety requirement is (that
indices are in bounds) as a comment
##########
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:
Something worth doing as a follow up prhaps
##########
arrow-select/src/coalesce/primitive.rs:
##########
@@ -229,6 +229,52 @@ impl<T: ArrowPrimitiveType + Debug> InProgressArray for
InProgressPrimitiveArray
}
}
+ unsafe fn copy_rows_by_indices(
Review Comment:
it would probably be good to document here in comments what is required for
safe usage of this function (eg that the indices are in bounds)
##########
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:
Is the idea here is to avoid adding multiple references to the same
underlying buffer on subsquent calls to take?
##########
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:
I wonder if we can avoid this check somehow (we would probably have to
vectorize the check in a first pass or something)
--
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]
