alamb commented on code in PR #11174:
URL: https://github.com/apache/arrow-rs/pull/11174#discussion_r4329746155
##########
arrow-select/src/coalesce.rs:
##########
@@ -336,11 +404,81 @@ impl BatchCoalescer {
batch: RecordBatch,
indices: &dyn Array,
) -> Result<(), ArrowError> {
- // todo: optimize this to avoid materializing (copying the results of
take indices to a new batch)
+ if let Some(u32_arr) = indices.as_any().downcast_ref::<UInt32Array>()
+ && u32_arr.null_count() == 0
+ {
+ validate_indices_bounds(u32_arr.values(), batch.num_rows())?;
Review Comment:
why only validate u32s?
Also this seems like it validates the indices, and then calls
take_record_batch which validates the indices again
##########
arrow-select/src/coalesce.rs:
##########
@@ -312,15 +377,18 @@ impl BatchCoalescer {
/// This is semantically equivalent of calling [`Self::push_batch`]
/// with the results from [`take_record_batch`]
///
+ /// Returns `Err(ArrowError::InvalidArgumentError)` if `indices` is a
+ /// non-null `UInt32Array` with any out-of-bounds value.
+ ///
/// # Example
/// ```
- /// # use arrow_array::{record_batch, UInt64Array};
+ /// # use arrow_array::{record_batch, UInt32Array};
Review Comment:
does this mean we can't use UInt64 as indexes anymore? According to Codex
they would stil work but not get the new fast path
##########
arrow-select/src/coalesce/byte_view.rs:
##########
@@ -354,13 +362,156 @@ 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,
+ ) {
+ 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) };
Review Comment:
I don't think the caller was required to do any validation (the comments say
they have to)
`debug_assert` isn't sufficient to ensure this invariant in release builds
-- we either need to actually verify it or mark the trait function as `unsafe`
##########
arrow-select/src/coalesce.rs:
##########
@@ -853,6 +991,23 @@ trait InProgressArray: std::fmt::Debug + Send + Sync {
}
}
+ /// Copy rows at `indices` from `source` into the in-progress array.
+ /// `indices` must already be bounds-checked against `source`.
Review Comment:
Update; see my comment above -- I think this actually needs to be marked
`unsafe`
##########
arrow-select/src/coalesce.rs:
##########
@@ -336,11 +404,81 @@ impl BatchCoalescer {
batch: RecordBatch,
indices: &dyn Array,
) -> Result<(), ArrowError> {
- // todo: optimize this to avoid materializing (copying the results of
take indices to a new batch)
+ if let Some(u32_arr) = indices.as_any().downcast_ref::<UInt32Array>()
+ && u32_arr.null_count() == 0
+ {
+ validate_indices_bounds(u32_arr.values(), batch.num_rows())?;
+ }
+
+ let exceeds_byte_view_threshold = self.has_byte_view_columns
+ && !should_use_byte_view_indexed_copy(batch.num_rows(),
indices.len());
+ if !should_use_indexed_copy(batch.num_rows(), indices.len()) ||
exceeds_byte_view_threshold
+ {
+ let taken_batch = take_record_batch(&batch, indices)?;
+ return self.push_batch(taken_batch);
+ }
+
+ let fast_path_allowed = self
+ .biggest_coalesce_batch_size
+ .is_none_or(|limit| indices.len() <= limit);
+ if fast_path_allowed
+ && let Some(u32_arr) =
indices.as_any().downcast_ref::<UInt32Array>()
+ && u32_arr.null_count() == 0
+ {
+ return self.push_batch_with_index(batch, u32_arr);
+ }
let taken_batch = take_record_batch(&batch, indices)?;
self.push_batch(taken_batch)
}
+ fn push_batch_with_index(
Review Comment:
could we commet that this is a fast path and under what circumstance sit is
allowed?
##########
arrow-select/src/coalesce.rs:
##########
@@ -853,6 +991,23 @@ trait InProgressArray: std::fmt::Debug + Send + Sync {
}
}
+ /// Copy rows at `indices` from `source` into the in-progress array.
+ /// `indices` must already be bounds-checked against `source`.
Review Comment:
I found this somewhat confusing as it says ``indices` must already be
bounds-checked against `source`.` but doesn't say what happens if this is not
the case (I was wondering if it could lead to out of bounds erorrs)
However, after review I think the result is that it panics
So I recommend adding a note here to clarify:
```suggestion
/// # Panics
/// Pancis if any value in `indices` is not a valid index in `source`.
##########
arrow-select/src/coalesce.rs:
##########
@@ -41,6 +42,49 @@ use fixed_size_binary::InProgressFixedSizeBinaryArray;
use generic::GenericInProgressArray;
use primitive::InProgressPrimitiveArray;
+// Caller must have bounds-checked `indices` against `source_nulls`.
Review Comment:
can we please document what happens if this is not the case?
If it results in an out of bounds access we need to mark this function as
`unsafe` -- a debug assert is not enough to ensure this
--
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]
