From 5c102e405dab4e4a704100f93f4b152d6a8ab694 Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Fri, 16 Oct 2026 11:36:44 -0400 Subject: [PATCH 1/6] perf(coalesce): directly copy sparse indexed rows Scatter sparse non-null UInt32 selections directly into specialized primitive and byte-view builders, avoiding the intermediate taken RecordBatch. Validate indices before builder mutation and keep dense selections on the existing take path. Compact sparse ByteView selections into bounded data blocks, while reusing source buffers for dense selections and repeated chunks. This keeps offsets valid when selected non-inline data exceeds a single Arrow buffer. Closes #8957 --- arrow-select/src/coalesce.rs | 275 ++++++++++++++++- arrow-select/src/coalesce/byte_view.rs | 398 ++++++++++++++++++++++++- arrow-select/src/coalesce/primitive.rs | 59 +++- arrow-select/src/take.rs | 2 +- 4 files changed, 724 insertions(+), 10 deletions(-) diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs index 6f5378a8f5c8..ecfa2e316a60 100644 --- a/arrow-select/src/coalesce.rs +++ b/arrow-select/src/coalesce.rs @@ -22,9 +22,10 @@ //! [`take`]: crate::take::take use crate::filter::{FilterBuilder, FilterPredicate, FilterSelection}; use crate::interleave::interleave_record_batch; -use crate::take::take_record_batch; -use arrow_array::types::{BinaryViewType, StringViewType}; -use arrow_array::{Array, ArrayRef, BooleanArray, RecordBatch, downcast_primitive}; +use crate::take::{take, take_bits, take_record_batch}; +use arrow_array::types::{BinaryViewType, StringViewType, UInt32Type}; +use arrow_array::{Array, ArrayRef, BooleanArray, RecordBatch, UInt32Array, downcast_primitive}; +use arrow_buffer::{Buffer, NullBuffer, NullBufferBuilder, ScalarBuffer}; use arrow_schema::{ArrowError, DataType, SchemaRef}; use std::collections::VecDeque; use std::sync::Arc; @@ -41,6 +42,45 @@ use fixed_size_binary::InProgressFixedSizeBinaryArray; use generic::GenericInProgressArray; use primitive::InProgressPrimitiveArray; +// Caller must have bounds-checked `indices` against `source_nulls`. +pub(crate) fn append_gathered_nulls( + nulls: &mut NullBufferBuilder, + source_nulls: &NullBuffer, + indices: &[u32], +) { + if source_nulls.null_count() == 0 { + nulls.append_n_non_nulls(indices.len()); + return; + } + let buf = Buffer::from_slice_ref(indices); + let scalar: ScalarBuffer = ScalarBuffer::new(buf, 0, indices.len()); + let indices_arr = UInt32Array::new(scalar, None); + let gathered = take_bits::(source_nulls.inner(), &indices_arr); + nulls.append_buffer(&NullBuffer::new(gathered)); +} + +pub(crate) fn validate_indices_bounds( + indices: &[u32], + source_len: usize, +) -> Result<(), ArrowError> { + let has_out_of_bounds = indices + .iter() + .copied() + .max() + .is_some_and(|idx| (idx as usize) >= source_len); + if has_out_of_bounds { + let bad = indices + .iter() + .copied() + .find(|&idx| (idx as usize) >= source_len) + .expect("maximum index was out of bounds"); + return Err(ArrowError::InvalidArgumentError(format!( + "take index {bad} out of bounds for array of {source_len} rows" + ))); + } + Ok(()) +} + fn has_sparse_filter_copy(data_type: &DataType) -> bool { data_type.is_primitive() || matches!(data_type, DataType::Utf8View | DataType::BinaryView) } @@ -56,6 +96,13 @@ fn should_use_sparse_filter_copy(filter_len: usize, selected_count: usize) -> bo selected_count <= filter_len / SPARSE_FILTER_COPY_MAX_SELECTIVITY_DENOMINATOR } +// Dense selections are faster through `take`. +const INDEXED_COPY_MAX_SELECTIVITY_DENOMINATOR: usize = 2; + +fn should_use_indexed_copy(input_rows: usize, selected_rows: usize) -> bool { + selected_rows <= input_rows / INDEXED_COPY_MAX_SELECTIVITY_DENOMINATOR +} + /// Concatenate multiple [`RecordBatch`]es /// /// Implements the common pattern of incrementally creating output @@ -336,11 +383,61 @@ 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 !should_use_indexed_copy(batch.num_rows(), indices.len()) { + let taken_batch = take_record_batch(&batch, indices)?; + return self.push_batch(taken_batch); + } + + // Let push_batch handle large-batch bypasses and unsupported columns. + let fast_path_allowed = self + .biggest_coalesce_batch_size + .is_none_or(|limit| indices.len() <= limit) + && !self.has_non_specialized_filter_columns; + if fast_path_allowed + && let Some(u32_arr) = indices.as_any().downcast_ref::() + && u32_arr.null_count() == 0 + { + return self.push_batch_with_index(batch, u32_arr.values()); + } let taken_batch = take_record_batch(&batch, indices)?; self.push_batch(taken_batch) } + fn push_batch_with_index( + &mut self, + batch: RecordBatch, + indices: &[u32], + ) -> Result<(), ArrowError> { + if indices.is_empty() { + return Ok(()); + } + let columns = batch.columns(); + if columns.len() != self.in_progress_arrays.len() { + return Err(ArrowError::InvalidArgumentError(format!( + "Batch has {} columns but BatchCoalescer expects {}", + columns.len(), + self.in_progress_arrays.len() + ))); + } + validate_indices_bounds(indices, batch.num_rows())?; + + let mut offset = 0; + while offset < indices.len() { + let room = self.target_batch_size - self.buffered_rows; + let take_n = (indices.len() - offset).min(room); + let chunk = &indices[offset..offset + take_n]; + for (in_progress, column) in self.in_progress_arrays.iter_mut().zip(columns) { + in_progress.copy_rows_by_indices(column, chunk)?; + } + self.buffered_rows += take_n; + offset += take_n; + if self.buffered_rows >= self.target_batch_size { + self.finish_buffered_batch()?; + } + } + Ok(()) + } + /// Push rows gathered from multiple `batches` into the Coalescer. /// /// Semantically equivalent to @@ -853,6 +950,26 @@ 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`. + fn copy_rows_by_indices( + &mut self, + source: &ArrayRef, + indices: &[u32], + ) -> Result<(), ArrowError> { + if indices.is_empty() { + return Ok(()); + } + let buf = Buffer::from_slice_ref(indices); + let scalar: ScalarBuffer = ScalarBuffer::new(buf, 0, indices.len()); + let indices_arr = UInt32Array::new(scalar, None); + let taken = take(source.as_ref(), &indices_arr, None)?; + self.set_source(Some(taken)); + let result = self.copy_rows(0, indices.len()); + self.set_source(None); + result + } + /// Finish the currently in-progress array and return it as an `ArrayRef` fn finish(&mut self) -> Result; @@ -865,6 +982,7 @@ mod tests { use super::*; use crate::concat::concat_batches; use crate::filter::filter_record_batch; + use crate::take::take_record_batch; use arrow_array::builder::StringViewBuilder; use arrow_array::cast::AsArray; use arrow_array::types::Int32Type; @@ -919,6 +1037,12 @@ mod tests { assert!(!should_use_sparse_filter_copy(8192, 6553)); } + #[test] + fn test_indexed_copy_threshold() { + assert!(should_use_indexed_copy(8192, 4096)); + assert!(!should_use_indexed_copy(8192, 4097)); + } + #[test] fn test_single_large_batch_greater_than_target() { // test a single large batch @@ -2988,4 +3112,147 @@ mod tests { "draining must release the accounted memory" ); } + + #[test] + fn test_push_batch_with_indices_splits_at_target_batch_size() { + let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])); + let array = Int32Array::from_iter_values(0..1000i32); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap(); + let indices = UInt32Array::from_iter_values((0..1000u32).step_by(3)); + let target_batch_size = 100; + + let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), target_batch_size); + coalescer + .push_batch_with_indices(batch.clone(), &indices) + .unwrap(); + coalescer.finish_buffered_batch().unwrap(); + let mut completed = vec![]; + while let Some(batch) = coalescer.next_completed_batch() { + completed.push(batch); + } + for batch in &completed[..completed.len() - 1] { + assert_eq!(batch.num_rows(), target_batch_size); + } + + let actual = concat_batches(&schema, completed.iter()).unwrap(); + let expected = take_record_batch(&batch, &indices).unwrap(); + assert_eq!(expected, actual); + } + + #[test] + fn test_push_batch_with_indices_empty_is_noop() { + let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])); + let array = Int32Array::from_iter_values(0..4i32); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap(); + let indices = UInt32Array::from(Vec::::new()); + + let mut coalescer = BatchCoalescer::new(schema, 4096); + coalescer.push_batch_with_indices(batch, &indices).unwrap(); + coalescer.finish_buffered_batch().unwrap(); + assert!(coalescer.next_completed_batch().is_none()); + } + + #[test] + fn test_push_batch_with_indices_rejects_schema_mismatch() { + let schema = Arc::new(Schema::new(vec![ + Field::new("c0", DataType::Int32, false), + Field::new("c1", DataType::Int32, false), + ])); + let mut coalescer = BatchCoalescer::new(schema, 100); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])), + vec![Arc::new(Int32Array::from_iter_values(0..5i32))], + ) + .unwrap(); + let indices = UInt32Array::from(vec![0u32, 1]); + let err = coalescer + .push_batch_with_indices(batch, &indices) + .unwrap_err() + .to_string(); + assert!( + err.contains("Batch has 1 columns but BatchCoalescer expects 2"), + "unexpected error: {err}" + ); + } + + #[test] + fn test_push_batch_with_indices_rejects_oob_before_mutation() { + let schema = Arc::new(Schema::new(vec![ + Field::new("c0", DataType::Int32, false), + Field::new("c1", DataType::Utf8View, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from_iter_values(0..16i32)), + Arc::new(StringViewArray::from(vec!["a"; 16])), + ], + ) + .unwrap(); + let indices = UInt32Array::from(vec![0u32, 1, 2, 3, 0, 1, 2, 99]); + let expected_err = "take index 99 out of bounds for array of 16 rows"; + + let mut coalescer = BatchCoalescer::new(schema, 4); + let err = coalescer + .push_batch_with_indices(batch, &indices) + .unwrap_err() + .to_string(); + assert!(err.contains(expected_err), "unexpected error: {err}"); + assert_eq!(coalescer.get_buffered_rows(), 0); + assert!(!coalescer.has_completed_batch()); + } + + #[test] + fn test_push_batch_with_indices_sparse_byte_view_bounded_size() { + let schema = Arc::new(Schema::new(vec![Field::new( + "v", + DataType::BinaryView, + false, + )])); + let rows_per_batch = 2048; + let num_batches = 64; + let target_batch_size = 4 * num_batches; + let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), target_batch_size); + + for batch_idx in 0..num_batches { + let values: Vec> = (0..rows_per_batch) + .map(|row| { + format!("batch={batch_idx:04} row={row:04} padding padding padding") + .into_bytes() + }) + .collect(); + let array = BinaryViewArray::from_iter(values.iter().map(|v| Some(v.as_slice()))); + assert!( + !array.data_buffers().is_empty(), + "need non-inline views to exercise the sparse path" + ); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap(); + let indices = UInt32Array::from(vec![0u32]); + coalescer.push_batch_with_indices(batch, &indices).unwrap(); + } + + let max_expected_size = 256 * 1024; + let retained_size = coalescer.size(); + assert!( + retained_size < max_expected_size, + "retained {retained_size} bytes; expected under {max_expected_size}", + ); + } + + #[test] + fn test_push_batch_with_indices_preserves_biggest_bypass() { + let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])); + let array = Int32Array::from_iter_values(0..1000i32); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap(); + + let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), 100) + .with_biggest_coalesce_batch_size(Some(50)); + + let indices = UInt32Array::from_iter_values(0..500u32); + coalescer.push_batch_with_indices(batch, &indices).unwrap(); + + let output = coalescer.next_completed_batch().expect("one batch"); + assert_eq!(output.num_rows(), 500); + assert!(coalescer.next_completed_batch().is_none()); + } } diff --git a/arrow-select/src/coalesce/byte_view.rs b/arrow-select/src/coalesce/byte_view.rs index a6697bfdb599..62d4b01c94a7 100644 --- a/arrow-select/src/coalesce/byte_view.rs +++ b/arrow-select/src/coalesce/byte_view.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use crate::coalesce::InProgressArray; +use crate::coalesce::{InProgressArray, append_gathered_nulls}; use crate::filter::{FilterPredicate, FilterSelection}; use arrow_array::cast::AsArray; use arrow_array::types::ByteViewType; @@ -24,7 +24,7 @@ use arrow_buffer::{Buffer, NullBuffer, NullBufferBuilder}; use arrow_data::{ByteView, MAX_INLINE_VIEW_LEN}; use arrow_schema::ArrowError; use std::marker::PhantomData; -use std::sync::Arc; +use std::sync::{Arc, Weak}; /// InProgressArray for [`StringViewArray`] and [`BinaryViewArray`] /// @@ -58,6 +58,12 @@ pub(crate) struct InProgressByteViewArray { 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>, + /// Position in `completed` where `last_indexed_source`'s buffers begin. + last_indexed_starting_buffer: u32, } struct Source { @@ -95,6 +101,8 @@ impl InProgressByteViewArray { completed: vec![], completed_buffers_size: 0, size_of_completed_buffers_from_current_source: 0, + last_indexed_source: None, + last_indexed_starting_buffer: 0, buffer_source, _phantom: PhantomData, } @@ -354,6 +362,149 @@ impl InProgressByteViewArray { 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::(); + 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); + } + copied_bytes = copied_bytes + .checked_add(str_len) + .expect("selected byte count overflow"); + } + // SAFETY: `reserve` above allocated room for every selected view. + unsafe { output.add(i).write(view.as_u128()) }; + } + // SAFETY: the loop above initialized exactly `indices.len()` views. + unsafe { self.views.set_len(current_len + indices.len()) }; + self.current = Some(dst_buffer); + } +} + +fn selected_non_inline_bytes( + src_views: &[u128], + indices: &[u32], + max_buffer_size: usize, +) -> Option { + let mut total = 0usize; + for &idx in indices { + // SAFETY: caller validated idx < src_views.len(). + let raw = unsafe { *src_views.get_unchecked(idx as usize) }; + let length = raw as u32; + if length > MAX_INLINE_VIEW_LEN { + let length = length as usize; + if length > max_buffer_size { + return None; + } + total = total.checked_add(length)?; + } + } + Some(total) } impl InProgressArray for InProgressByteViewArray { @@ -361,6 +512,7 @@ impl InProgressArray for InProgressByteViewArray { // If used values from source, add only the size that was used self.completed_buffers_size += self.size_of_completed_buffers_from_current_source; self.size_of_completed_buffers_from_current_source = 0; + self.last_indexed_source = None; self.source = source.map(|array| { let s = array.as_byte_view::(); @@ -487,6 +639,114 @@ impl InProgressArray for InProgressByteViewArray { Ok(()) } + fn copy_rows_by_indices( + &mut self, + source: &ArrayRef, + indices: &[u32], + ) -> Result<(), ArrowError> { + if indices.is_empty() { + return Ok(()); + } + let typed = source.as_byte_view::(); + self.ensure_capacity(); + + if let Some(src_nulls) = typed.nulls() { + append_gathered_nulls(&mut self.nulls, src_nulls, indices); + } else { + self.nulls.append_n_non_nulls(indices.len()); + } + + let src_views = typed.views(); + let src_buffers = typed.data_buffers(); + + if src_buffers.is_empty() { + let current_len = self.views.len(); + self.views.reserve(indices.len()); + // SAFETY: caller validated indices < typed.len() == src_views.len(). + unsafe { + let base = self.views.spare_capacity_mut().as_mut_ptr().cast::(); + let src = src_views.as_ptr(); + for (i, &idx) in indices.iter().enumerate() { + base.add(i).write(*src.add(idx as usize)); + } + self.views.set_len(current_len + indices.len()); + } + return Ok(()); + } + + let same_source = self + .last_indexed_source + .as_ref() + .is_some_and(|cached| std::ptr::addr_eq(cached.as_ptr(), Arc::as_ptr(source))); + + if !same_source + && let Some(selected_bytes) = + selected_non_inline_bytes(src_views, indices, i32::MAX as usize) + { + let source_capacity = src_buffers + .iter() + .try_fold(0usize, |total, buffer| total.checked_add(buffer.capacity())); + if let (Some(source_capacity), Some(compaction_threshold)) = + (source_capacity, selected_bytes.checked_mul(2)) + && source_capacity > compaction_threshold + { + self.append_indexed_views_and_copy_strings( + src_views, + indices, + src_buffers, + selected_bytes, + ); + return Ok(()); + } + } + + let starting_buffer: u32 = if same_source { + self.last_indexed_starting_buffer + } else { + if let Some(buffer) = self.current.take() { + let buffer: Buffer = buffer.into(); + self.completed_buffers_size += buffer.capacity(); + self.completed.push(buffer); + } + let start: u32 = self.completed.len().try_into().expect("too many buffers"); + let buffers_size = src_buffers.iter().map(|b| b.capacity()).sum::(); + self.completed_buffers_size += buffers_size; + self.completed.extend_from_slice(src_buffers); + self.last_indexed_source = Some(Arc::downgrade(source)); + self.last_indexed_starting_buffer = start; + start + }; + + let current_len = self.views.len(); + self.views.reserve(indices.len()); + // SAFETY: caller validated indices < typed.len() == src_views.len(). + unsafe { + let base = self.views.spare_capacity_mut().as_mut_ptr().cast::(); + let src = src_views.as_ptr(); + if starting_buffer == 0 { + for (i, &idx) in indices.iter().enumerate() { + base.add(i).write(*src.add(idx as usize)); + } + } else { + // Non-inline views (length > MAX_INLINE_VIEW_LEN) carry buffer_index + // in bits [64, 96); bump via raw u128 add to skip the ByteView round-trip. + let bump = (starting_buffer as u128) << 64; + for (i, &idx) in indices.iter().enumerate() { + let raw = *src.add(idx as usize); + let len = raw as u32; + let adjusted = if len > MAX_INLINE_VIEW_LEN { + raw + bump + } else { + raw + }; + base.add(i).write(adjusted); + } + } + self.views.set_len(current_len + indices.len()); + } + Ok(()) + } + fn finish(&mut self) -> Result { self.finish_current(); assert!(self.current.is_none()); @@ -498,6 +758,8 @@ impl InProgressArray for InProgressByteViewArray { // Not reusing anything since we took all complete self.size_of_completed_buffers_from_current_source = 0; self.completed_buffers_size = 0; + self.last_indexed_source = None; + self.last_indexed_starting_buffer = 0; // Safety: we created valid views and buffers above and the // input arrays had value data and nulls @@ -562,9 +824,11 @@ impl BufferSource { #[cfg(test)] mod tests { use super::*; + use crate::concat::concat; use crate::filter::FilterBuilder; + use crate::take::take; use arrow_array::types::BinaryViewType; - use arrow_array::{BinaryViewArray, BooleanArray}; + use arrow_array::{BinaryViewArray, BooleanArray, UInt32Array}; #[test] fn test_buffer_source() { @@ -737,4 +1001,132 @@ mod tests { assert_eq!(in_progress_size_with_split, in_progress_size_without_split); } + + #[test] + fn test_copy_rows_by_indices_inline_only() { + let values: Vec> = vec![ + Some(b"a"), + Some(b"bb"), + None, + Some(b"ccc"), + Some(b"dddddddddddd"), + Some(b"eee"), + ]; + let source: ArrayRef = Arc::new(BinaryViewArray::from_iter(values.iter().copied())); + assert!(source.as_binary_view().data_buffers().is_empty()); + let indices = [5u32, 0, 2, 4, 1, 1, 3, 0]; + + let mut in_progress = InProgressByteViewArray::::new(16); + in_progress.copy_rows_by_indices(&source, &indices).unwrap(); + let actual = in_progress.finish().unwrap(); + + let expected = take(source.as_ref(), &UInt32Array::from(indices.to_vec()), None).unwrap(); + assert_eq!(actual.as_ref(), expected.as_ref()); + assert!(actual.as_binary_view().data_buffers().is_empty()); + } + + #[test] + fn test_copy_rows_by_indices_reuses_same_source_buffers() { + let (array, _) = non_inline_array(32); + let source: ArrayRef = Arc::new(array); + + let mut in_progress = InProgressByteViewArray::::new(64); + + let dense_indices: Vec = (0..32u32).collect(); + in_progress + .copy_rows_by_indices(&source, &dense_indices) + .unwrap(); + let after_first = in_progress.size(); + + in_progress + .copy_rows_by_indices(&source, &[1u32, 2, 4]) + .unwrap(); + let after_second = in_progress.size(); + + let second_growth = after_second - after_first; + assert!( + second_growth <= 3 * size_of::(), + "same-source re-entry added {second_growth} bytes" + ); + + in_progress.finish().unwrap(); + } + + #[test] + fn test_copy_rows_by_indices_remaps_buffer_index() { + let (array_a, _) = non_inline_array(8); + let (array_b, _) = non_inline_array(8); + let source_a: ArrayRef = Arc::new(array_a); + let source_b: ArrayRef = Arc::new(array_b); + + let indices_a = [0u32, 1, 3, 5, 7]; + let indices_b = [2u32, 5, 1, 6, 4]; + + let mut in_progress = InProgressByteViewArray::::new(16); + in_progress + .copy_rows_by_indices(&source_a, &indices_a) + .unwrap(); + in_progress + .copy_rows_by_indices(&source_b, &indices_b) + .unwrap(); + let actual = in_progress.finish().unwrap(); + + let taken_a = take( + source_a.as_ref(), + &UInt32Array::from(indices_a.to_vec()), + None, + ) + .unwrap(); + let taken_b = take( + source_b.as_ref(), + &UInt32Array::from(indices_b.to_vec()), + None, + ) + .unwrap(); + let expected = concat(&[taken_a.as_ref(), taken_b.as_ref()]).unwrap(); + assert_eq!(actual.as_ref(), expected.as_ref()); + + let actual = actual.as_binary_view(); + assert_eq!(actual.data_buffers().len(), 2); + for raw in &actual.views()[..indices_a.len()] { + assert_eq!(ByteView::from(*raw).buffer_index, 0); + } + for raw in &actual.views()[indices_a.len()..] { + assert_eq!(ByteView::from(*raw).buffer_index, 1); + } + } + + #[test] + fn test_copy_rows_by_indices_splits_compacted_buffers() { + let (array, _) = non_inline_array(8); + let source: ArrayRef = Arc::new(array); + let indices = [0u32, 3, 7]; + let max_buffer_size = 80; + let typed = source.as_binary_view(); + let selected_bytes = + selected_non_inline_bytes(typed.views(), &indices, max_buffer_size).unwrap(); + + let mut in_progress = InProgressByteViewArray::::new(16); + in_progress.nulls.append_n_non_nulls(indices.len()); + in_progress.append_indexed_views_and_copy_strings_with_max_buffer_size( + typed.views(), + &indices, + typed.data_buffers(), + selected_bytes, + max_buffer_size, + ); + let actual = in_progress.finish().unwrap(); + + let expected = take(source.as_ref(), &UInt32Array::from(indices.to_vec()), None).unwrap(); + assert_eq!(actual.as_ref(), expected.as_ref()); + assert!(actual.as_binary_view().data_buffers().len() > 1); + assert!( + actual + .as_binary_view() + .data_buffers() + .iter() + .all(|buffer| buffer.len() <= max_buffer_size) + ); + actual.to_data().validate_full().unwrap(); + } } diff --git a/arrow-select/src/coalesce/primitive.rs b/arrow-select/src/coalesce/primitive.rs index 62460ab677bd..a2b0f537d860 100644 --- a/arrow-select/src/coalesce/primitive.rs +++ b/arrow-select/src/coalesce/primitive.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use crate::coalesce::InProgressArray; +use crate::coalesce::{InProgressArray, append_gathered_nulls}; use crate::filter::{FilterIndices, FilterPredicate, FilterSelection, FilterSlices}; use arrow_array::cast::AsArray; use arrow_array::{Array, ArrayRef, ArrowPrimitiveType, PrimitiveArray}; @@ -206,6 +206,43 @@ impl InProgressArray for InProgressPrimitiveArray } } + fn copy_rows_by_indices( + &mut self, + source: &ArrayRef, + indices: &[u32], + ) -> Result<(), ArrowError> { + if indices.is_empty() { + return Ok(()); + } + let typed = source.as_primitive::(); + self.ensure_capacity(); + + if let Some(src_nulls) = typed.nulls() { + append_gathered_nulls(&mut self.nulls, src_nulls, indices); + } else { + self.nulls.append_n_non_nulls(indices.len()); + } + + let values = typed.values(); + let current_len = self.current.len(); + self.current.reserve(indices.len()); + // SAFETY: indices were bounds-checked by the coalescer, and reserve + // allocated one output slot per index. + unsafe { + let base = self + .current + .spare_capacity_mut() + .as_mut_ptr() + .cast::(); + let src = values.as_ptr(); + for (i, &idx) in indices.iter().enumerate() { + base.add(i).write(*src.add(idx as usize)); + } + self.current.set_len(current_len + indices.len()); + } + Ok(()) + } + fn finish(&mut self) -> Result { // take and reset the current values and nulls let values = std::mem::take(&mut self.current); @@ -231,8 +268,9 @@ impl InProgressArray for InProgressPrimitiveArray mod tests { use super::*; use crate::filter::FilterBuilder; + use crate::take::take; use arrow_array::types::Int32Type; - use arrow_array::{BooleanArray, Int32Array}; + use arrow_array::{BooleanArray, Int32Array, UInt32Array}; #[test] fn test_copy_rows_by_filter_index_iterator() { @@ -373,4 +411,21 @@ mod tests { "null buffer must be included in size(): with_nulls={with_nulls} all_valid={all_valid}" ); } + + #[test] + fn test_copy_rows_by_indices() { + let source: ArrayRef = Arc::new(Int32Array::from_iter( + (0..32i32).map(|i| if i % 3 == 0 { None } else { Some(i) }), + )); + let indices = [ + 0u32, 15, 2, 7, 3, 9, 12, 1, 20, 15, 6, 8, 11, 18, 21, 24, 30, + ]; + + let mut in_progress = InProgressPrimitiveArray::::new(32, DataType::Int32); + in_progress.copy_rows_by_indices(&source, &indices).unwrap(); + let actual = in_progress.finish().unwrap(); + + let expected = take(source.as_ref(), &UInt32Array::from(indices.to_vec()), None).unwrap(); + assert_eq!(actual.as_ref(), expected.as_ref()); + } } diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index e4e5f534bf9a..1707301aa759 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -657,7 +657,7 @@ unsafe fn pack_bit(src: *const u8, bit_idx: usize, out_pos: usize) -> u8 { } #[inline(never)] -fn take_bits( +pub(crate) fn take_bits( values: &BooleanBuffer, indices: &PrimitiveArray, ) -> BooleanBuffer { From 0c012ed861d1999f8e9d1e23a0f02b1906fb2f6f Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Mon, 19 Oct 2026 14:17:34 -0400 Subject: [PATCH 2/6] perf(coalesce): cap byte-view indexed copy density Keep the existing primitive cutoff, but use the materialized take path above 1/16 selectivity when a batch contains Utf8View or BinaryView columns. --- arrow-select/src/coalesce.rs | 33 ++++++++++++++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs index ecfa2e316a60..4562bfe3ccbf 100644 --- a/arrow-select/src/coalesce.rs +++ b/arrow-select/src/coalesce.rs @@ -99,10 +99,17 @@ fn should_use_sparse_filter_copy(filter_len: usize, selected_count: usize) -> bo // Dense selections are faster through `take`. const INDEXED_COPY_MAX_SELECTIVITY_DENOMINATOR: usize = 2; +// Byte-view indexed copying is only faster for sparse selections. +const BYTE_VIEW_INDEXED_COPY_MAX_SELECTIVITY_DENOMINATOR: usize = 16; + fn should_use_indexed_copy(input_rows: usize, selected_rows: usize) -> bool { selected_rows <= input_rows / INDEXED_COPY_MAX_SELECTIVITY_DENOMINATOR } +fn should_use_byte_view_indexed_copy(input_rows: usize, selected_rows: usize) -> bool { + selected_rows <= input_rows / BYTE_VIEW_INDEXED_COPY_MAX_SELECTIVITY_DENOMINATOR +} + /// Concatenate multiple [`RecordBatch`]es /// /// Implements the common pattern of incrementally creating output @@ -206,6 +213,8 @@ pub struct BatchCoalescer { in_progress_arrays: Vec>, /// True if some column still needs the materialized filter path. has_non_specialized_filter_columns: bool, + /// True if some column uses the byte-view indexed-copy path. + has_byte_view_columns: bool, /// Buffered row count. Always less than `batch_size` buffered_rows: usize, /// Completed batches @@ -227,6 +236,10 @@ impl BatchCoalescer { .fields() .iter() .any(|field| !has_sparse_filter_copy(field.data_type())); + let has_byte_view_columns = schema + .fields() + .iter() + .any(|field| matches!(field.data_type(), DataType::Utf8View | DataType::BinaryView)); let in_progress_arrays = schema .fields() .iter() @@ -238,6 +251,7 @@ impl BatchCoalescer { target_batch_size, in_progress_arrays, has_non_specialized_filter_columns, + has_byte_view_columns, // We will for sure store at least one completed batch completed: VecDeque::with_capacity(1), buffered_rows: 0, @@ -383,7 +397,17 @@ impl BatchCoalescer { batch: RecordBatch, indices: &dyn Array, ) -> Result<(), ArrowError> { - if !should_use_indexed_copy(batch.num_rows(), indices.len()) { + let exceeds_byte_view_threshold = self.has_byte_view_columns + && !should_use_byte_view_indexed_copy(batch.num_rows(), indices.len()); + if exceeds_byte_view_threshold + && let Some(u32_arr) = indices.as_any().downcast_ref::() + && u32_arr.null_count() == 0 + { + validate_indices_bounds(u32_arr.values(), batch.num_rows())?; + } + + 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); } @@ -1043,6 +1067,13 @@ mod tests { assert!(!should_use_indexed_copy(8192, 4097)); } + #[test] + fn test_byte_view_indexed_copy_threshold() { + assert!(should_use_byte_view_indexed_copy(8192, 512)); + assert!(!should_use_byte_view_indexed_copy(8192, 513)); + assert!(!should_use_byte_view_indexed_copy(8192, 819)); + } + #[test] fn test_single_large_batch_greater_than_target() { // test a single large batch From e31420bf3b6f3959e2a083bfdfa7dc68376580c3 Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Tue, 20 Oct 2026 09:10:28 -0400 Subject: [PATCH 3/6] Addressed review feedback Hoist UInt32 OOB validation to the public entry so the error is uniform across copy paths, drop the double validation, and keep the InProgressArray trait on UInt32Array with the per-chunk alloc removed. --- arrow-select/src/coalesce.rs | 134 ++++++++++++++++++++----- arrow-select/src/coalesce/byte_view.rs | 54 ++++------ arrow-select/src/coalesce/primitive.rs | 13 +-- 3 files changed, 135 insertions(+), 66 deletions(-) diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs index 4562bfe3ccbf..22cca83f742f 100644 --- a/arrow-select/src/coalesce.rs +++ b/arrow-select/src/coalesce.rs @@ -25,7 +25,7 @@ use crate::interleave::interleave_record_batch; use crate::take::{take, take_bits, take_record_batch}; use arrow_array::types::{BinaryViewType, StringViewType, UInt32Type}; use arrow_array::{Array, ArrayRef, BooleanArray, RecordBatch, UInt32Array, downcast_primitive}; -use arrow_buffer::{Buffer, NullBuffer, NullBufferBuilder, ScalarBuffer}; +use arrow_buffer::{NullBuffer, NullBufferBuilder}; use arrow_schema::{ArrowError, DataType, SchemaRef}; use std::collections::VecDeque; use std::sync::Arc; @@ -46,16 +46,20 @@ use primitive::InProgressPrimitiveArray; pub(crate) fn append_gathered_nulls( nulls: &mut NullBufferBuilder, source_nulls: &NullBuffer, - indices: &[u32], + 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 buf = Buffer::from_slice_ref(indices); - let scalar: ScalarBuffer = ScalarBuffer::new(buf, 0, indices.len()); - let indices_arr = UInt32Array::new(scalar, None); - let gathered = take_bits::(source_nulls.inner(), &indices_arr); + let gathered = take_bits::(source_nulls.inner(), indices); nulls.append_buffer(&NullBuffer::new(gathered)); } @@ -373,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}; /// # use arrow_select::coalesce::BatchCoalescer; /// let batch1 = record_batch!(("a", Int32, [0, 0, 0])).unwrap(); /// let batch2 = record_batch!(("a", Int32, [1, 1, 4, 5, 1, 4])).unwrap(); /// // Sorted indices to create a sorted output, this can be obtained with /// // `arrow-ord`'s sort_to_indices operation - /// let indices = UInt64Array::from(vec![0, 1, 4, 2, 5, 3]); + /// let indices = UInt32Array::from(vec![0, 1, 4, 2, 5, 3]); /// // create a new Coalescer that targets creating 1000 row batches /// let mut coalescer = BatchCoalescer::new(batch1.schema(), 1000); /// coalescer.push_batch(batch1); @@ -397,31 +404,28 @@ impl BatchCoalescer { batch: RecordBatch, indices: &dyn Array, ) -> Result<(), ArrowError> { - let exceeds_byte_view_threshold = self.has_byte_view_columns - && !should_use_byte_view_indexed_copy(batch.num_rows(), indices.len()); - if exceeds_byte_view_threshold - && let Some(u32_arr) = indices.as_any().downcast_ref::() + if let Some(u32_arr) = indices.as_any().downcast_ref::() && 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 push_batch handle large-batch bypasses and unsupported columns. let fast_path_allowed = self .biggest_coalesce_batch_size - .is_none_or(|limit| indices.len() <= limit) - && !self.has_non_specialized_filter_columns; + .is_none_or(|limit| indices.len() <= limit); if fast_path_allowed && let Some(u32_arr) = indices.as_any().downcast_ref::() && u32_arr.null_count() == 0 { - return self.push_batch_with_index(batch, u32_arr.values()); + return self.push_batch_with_index(batch, u32_arr); } let taken_batch = take_record_batch(&batch, indices)?; self.push_batch(taken_batch) @@ -430,7 +434,7 @@ impl BatchCoalescer { fn push_batch_with_index( &mut self, batch: RecordBatch, - indices: &[u32], + indices: &UInt32Array, ) -> Result<(), ArrowError> { if indices.is_empty() { return Ok(()); @@ -443,13 +447,26 @@ impl BatchCoalescer { self.in_progress_arrays.len() ))); } - validate_indices_bounds(indices, batch.num_rows())?; + debug_assert!( + indices.null_count() == 0 + && indices + .values() + .iter() + .all(|&idx| (idx as usize) < batch.num_rows()), + "push_batch_with_index requires caller-validated non-null, in-bounds indices", + ); let mut offset = 0; while offset < indices.len() { let room = self.target_batch_size - self.buffered_rows; let take_n = (indices.len() - offset).min(room); - let chunk = &indices[offset..offset + take_n]; + let sliced; + let chunk: &UInt32Array = if offset == 0 && take_n == indices.len() { + indices + } else { + sliced = indices.slice(offset, take_n); + &sliced + }; for (in_progress, column) in self.in_progress_arrays.iter_mut().zip(columns) { in_progress.copy_rows_by_indices(column, chunk)?; } @@ -979,15 +996,12 @@ trait InProgressArray: std::fmt::Debug + Send + Sync { fn copy_rows_by_indices( &mut self, source: &ArrayRef, - indices: &[u32], + indices: &UInt32Array, ) -> Result<(), ArrowError> { if indices.is_empty() { return Ok(()); } - let buf = Buffer::from_slice_ref(indices); - let scalar: ScalarBuffer = ScalarBuffer::new(buf, 0, indices.len()); - let indices_arr = UInt32Array::new(scalar, None); - let taken = take(source.as_ref(), &indices_arr, None)?; + let taken = take(source.as_ref(), indices, None)?; self.set_source(Some(taken)); let result = self.copy_rows(0, indices.len()); self.set_source(None); @@ -3233,6 +3247,78 @@ mod tests { assert!(!coalescer.has_completed_batch()); } + #[test] + fn test_push_batch_with_indices_rejects_oob_in_direct_copy_chunk() { + let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from_iter_values(0..64i32))], + ) + .unwrap(); + let indices = UInt32Array::from(vec![0u32, 1, 99]); + + let mut coalescer = BatchCoalescer::new(schema, 2); + let err = coalescer + .push_batch_with_indices(batch, &indices) + .unwrap_err() + .to_string(); + assert!( + err.contains("take index 99 out of bounds for array of 64 rows"), + "unexpected error: {err}" + ); + assert_eq!(coalescer.get_buffered_rows(), 0); + assert!(!coalescer.has_completed_batch()); + } + + #[test] + fn test_push_batch_with_indices_oob_error_is_uniform_across_schemas() { + let int_only_schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])); + let int_batch = RecordBatch::try_new( + Arc::clone(&int_only_schema), + vec![Arc::new(Int32Array::from_iter_values(0..64i32))], + ) + .unwrap(); + + let view_schema = Arc::new(Schema::new(vec![Field::new( + "c0", + DataType::Utf8View, + false, + )])); + let view_batch = RecordBatch::try_new( + Arc::clone(&view_schema), + vec![Arc::new(StringViewArray::from(vec!["a"; 64]))], + ) + .unwrap(); + + let short_indices = UInt32Array::from(vec![0u32, 1, 99]); + let long_indices = { + let mut values: Vec = (0..39).collect(); + values.push(99); + UInt32Array::from(values) + }; + + let expected_err = "take index 99 out of bounds for array of 64 rows"; + for (schema, batch, indices) in [ + ( + Arc::clone(&int_only_schema), + int_batch.clone(), + &short_indices, + ), + (Arc::clone(&int_only_schema), int_batch, &long_indices), + (Arc::clone(&view_schema), view_batch.clone(), &short_indices), + (Arc::clone(&view_schema), view_batch, &long_indices), + ] { + let mut coalescer = BatchCoalescer::new(schema, 100); + let err = coalescer + .push_batch_with_indices(batch, indices) + .unwrap_err() + .to_string(); + assert!(err.contains(expected_err), "unexpected error: {err}"); + assert_eq!(coalescer.get_buffered_rows(), 0); + assert!(!coalescer.has_completed_batch()); + } + } + #[test] fn test_push_batch_with_indices_sparse_byte_view_bounded_size() { let schema = Arc::new(Schema::new(vec![Field::new( diff --git a/arrow-select/src/coalesce/byte_view.rs b/arrow-select/src/coalesce/byte_view.rs index 62d4b01c94a7..f71e5616c9fa 100644 --- a/arrow-select/src/coalesce/byte_view.rs +++ b/arrow-select/src/coalesce/byte_view.rs @@ -19,7 +19,7 @@ use crate::coalesce::{InProgressArray, append_gathered_nulls}; use crate::filter::{FilterPredicate, FilterSelection}; use arrow_array::cast::AsArray; use arrow_array::types::ByteViewType; -use arrow_array::{Array, ArrayRef, GenericByteViewArray}; +use arrow_array::{Array, ArrayRef, GenericByteViewArray, UInt32Array}; use arrow_buffer::{Buffer, NullBuffer, NullBufferBuilder}; use arrow_data::{ByteView, MAX_INLINE_VIEW_LEN}; use arrow_schema::ArrowError; @@ -462,17 +462,8 @@ impl InProgressByteViewArray { .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); - } + debug_assert!(dst_buffer.capacity() - dst_buffer.len() >= str_len); + dst_buffer.extend_from_slice(src); copied_bytes = copied_bytes .checked_add(str_len) .expect("selected byte count overflow"); @@ -642,7 +633,7 @@ impl InProgressArray for InProgressByteViewArray { fn copy_rows_by_indices( &mut self, source: &ArrayRef, - indices: &[u32], + indices: &UInt32Array, ) -> Result<(), ArrowError> { if indices.is_empty() { return Ok(()); @@ -656,6 +647,7 @@ impl InProgressArray for InProgressByteViewArray { self.nulls.append_n_non_nulls(indices.len()); } + let index_values = indices.values(); let src_views = typed.views(); let src_buffers = typed.data_buffers(); @@ -666,7 +658,7 @@ impl InProgressArray for InProgressByteViewArray { unsafe { let base = self.views.spare_capacity_mut().as_mut_ptr().cast::(); let src = src_views.as_ptr(); - for (i, &idx) in indices.iter().enumerate() { + for (i, &idx) in index_values.iter().enumerate() { base.add(i).write(*src.add(idx as usize)); } self.views.set_len(current_len + indices.len()); @@ -681,7 +673,7 @@ impl InProgressArray for InProgressByteViewArray { if !same_source && let Some(selected_bytes) = - selected_non_inline_bytes(src_views, indices, i32::MAX as usize) + selected_non_inline_bytes(src_views, index_values, i32::MAX as usize) { let source_capacity = src_buffers .iter() @@ -692,7 +684,7 @@ impl InProgressArray for InProgressByteViewArray { { self.append_indexed_views_and_copy_strings( src_views, - indices, + index_values, src_buffers, selected_bytes, ); @@ -724,14 +716,14 @@ impl InProgressArray for InProgressByteViewArray { let base = self.views.spare_capacity_mut().as_mut_ptr().cast::(); let src = src_views.as_ptr(); if starting_buffer == 0 { - for (i, &idx) in indices.iter().enumerate() { + for (i, &idx) in index_values.iter().enumerate() { base.add(i).write(*src.add(idx as usize)); } } else { // Non-inline views (length > MAX_INLINE_VIEW_LEN) carry buffer_index // in bits [64, 96); bump via raw u128 add to skip the ByteView round-trip. let bump = (starting_buffer as u128) << 64; - for (i, &idx) in indices.iter().enumerate() { + for (i, &idx) in index_values.iter().enumerate() { let raw = *src.add(idx as usize); let len = raw as u32; let adjusted = if len > MAX_INLINE_VIEW_LEN { @@ -1014,13 +1006,13 @@ mod tests { ]; let source: ArrayRef = Arc::new(BinaryViewArray::from_iter(values.iter().copied())); assert!(source.as_binary_view().data_buffers().is_empty()); - let indices = [5u32, 0, 2, 4, 1, 1, 3, 0]; + let indices = UInt32Array::from(vec![5u32, 0, 2, 4, 1, 1, 3, 0]); let mut in_progress = InProgressByteViewArray::::new(16); in_progress.copy_rows_by_indices(&source, &indices).unwrap(); let actual = in_progress.finish().unwrap(); - let expected = take(source.as_ref(), &UInt32Array::from(indices.to_vec()), None).unwrap(); + let expected = take(source.as_ref(), &indices, None).unwrap(); assert_eq!(actual.as_ref(), expected.as_ref()); assert!(actual.as_binary_view().data_buffers().is_empty()); } @@ -1032,14 +1024,14 @@ mod tests { let mut in_progress = InProgressByteViewArray::::new(64); - let dense_indices: Vec = (0..32u32).collect(); + let dense_indices = UInt32Array::from_iter_values(0..32u32); in_progress .copy_rows_by_indices(&source, &dense_indices) .unwrap(); let after_first = in_progress.size(); in_progress - .copy_rows_by_indices(&source, &[1u32, 2, 4]) + .copy_rows_by_indices(&source, &UInt32Array::from(vec![1u32, 2, 4])) .unwrap(); let after_second = in_progress.size(); @@ -1059,8 +1051,8 @@ mod tests { let source_a: ArrayRef = Arc::new(array_a); let source_b: ArrayRef = Arc::new(array_b); - let indices_a = [0u32, 1, 3, 5, 7]; - let indices_b = [2u32, 5, 1, 6, 4]; + let indices_a = UInt32Array::from(vec![0u32, 1, 3, 5, 7]); + let indices_b = UInt32Array::from(vec![2u32, 5, 1, 6, 4]); let mut in_progress = InProgressByteViewArray::::new(16); in_progress @@ -1071,18 +1063,8 @@ mod tests { .unwrap(); let actual = in_progress.finish().unwrap(); - let taken_a = take( - source_a.as_ref(), - &UInt32Array::from(indices_a.to_vec()), - None, - ) - .unwrap(); - let taken_b = take( - source_b.as_ref(), - &UInt32Array::from(indices_b.to_vec()), - None, - ) - .unwrap(); + let taken_a = take(source_a.as_ref(), &indices_a, None).unwrap(); + let taken_b = take(source_b.as_ref(), &indices_b, None).unwrap(); let expected = concat(&[taken_a.as_ref(), taken_b.as_ref()]).unwrap(); assert_eq!(actual.as_ref(), expected.as_ref()); diff --git a/arrow-select/src/coalesce/primitive.rs b/arrow-select/src/coalesce/primitive.rs index a2b0f537d860..005286aa6ce9 100644 --- a/arrow-select/src/coalesce/primitive.rs +++ b/arrow-select/src/coalesce/primitive.rs @@ -18,7 +18,7 @@ use crate::coalesce::{InProgressArray, append_gathered_nulls}; use crate::filter::{FilterIndices, FilterPredicate, FilterSelection, FilterSlices}; use arrow_array::cast::AsArray; -use arrow_array::{Array, ArrayRef, ArrowPrimitiveType, PrimitiveArray}; +use arrow_array::{Array, ArrayRef, ArrowPrimitiveType, PrimitiveArray, UInt32Array}; use arrow_buffer::{NullBuffer, NullBufferBuilder, ScalarBuffer}; use arrow_schema::{ArrowError, DataType}; use std::fmt::Debug; @@ -209,7 +209,7 @@ impl InProgressArray for InProgressPrimitiveArray fn copy_rows_by_indices( &mut self, source: &ArrayRef, - indices: &[u32], + indices: &UInt32Array, ) -> Result<(), ArrowError> { if indices.is_empty() { return Ok(()); @@ -223,6 +223,7 @@ impl InProgressArray for InProgressPrimitiveArray 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()); @@ -235,7 +236,7 @@ impl InProgressArray for InProgressPrimitiveArray .as_mut_ptr() .cast::(); let src = values.as_ptr(); - for (i, &idx) in indices.iter().enumerate() { + for (i, &idx) in index_values.iter().enumerate() { base.add(i).write(*src.add(idx as usize)); } self.current.set_len(current_len + indices.len()); @@ -417,15 +418,15 @@ mod tests { let source: ArrayRef = Arc::new(Int32Array::from_iter( (0..32i32).map(|i| if i % 3 == 0 { None } else { Some(i) }), )); - let indices = [ + let indices = UInt32Array::from(vec![ 0u32, 15, 2, 7, 3, 9, 12, 1, 20, 15, 6, 8, 11, 18, 21, 24, 30, - ]; + ]); let mut in_progress = InProgressPrimitiveArray::::new(32, DataType::Int32); in_progress.copy_rows_by_indices(&source, &indices).unwrap(); let actual = in_progress.finish().unwrap(); - let expected = take(source.as_ref(), &UInt32Array::from(indices.to_vec()), None).unwrap(); + let expected = take(source.as_ref(), &indices, None).unwrap(); assert_eq!(actual.as_ref(), expected.as_ref()); } } From aaa05b39e1150db732b37b0e4d5f8a6e860367d5 Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Tue, 20 Oct 2026 10:07:08 -0400 Subject: [PATCH 4/6] add test that covers non-special fast path types --- arrow-select/src/coalesce.rs | 82 +++++++++++++++++++++++++++++++++++- 1 file changed, 80 insertions(+), 2 deletions(-) diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs index 22cca83f742f..cfd78b6cdfff 100644 --- a/arrow-select/src/coalesce.rs +++ b/arrow-select/src/coalesce.rs @@ -1025,8 +1025,9 @@ mod tests { use arrow_array::cast::AsArray; use arrow_array::types::Int32Type; use arrow_array::{ - BinaryViewArray, Int32Array, Int64Array, RecordBatchOptions, StringArray, StringViewArray, - TimestampNanosecondArray, UInt32Array, UInt64Array, make_array, record_batch, + BinaryViewArray, FixedSizeBinaryArray, Int32Array, Int64Array, RecordBatchOptions, + StringArray, StringViewArray, TimestampNanosecondArray, UInt32Array, UInt64Array, + make_array, record_batch, }; use arrow_buffer::BooleanBufferBuilder; use arrow_schema::{DataType, Field, Schema}; @@ -3184,6 +3185,83 @@ mod tests { assert_eq!(expected, actual); } + #[test] + fn test_push_batch_with_indices_mixed_schema_with_nulls_splits_matches_take() { + let fsb_size: i32 = 4; + let schema = Arc::new(Schema::new(vec![ + Field::new("i", DataType::Int32, true), + Field::new("s", DataType::Utf8, true), + Field::new("b", DataType::Boolean, true), + Field::new("f", DataType::FixedSizeBinary(fsb_size), true), + ])); + + let num_rows = 300usize; + let int_col = Int32Array::from_iter( + (0..num_rows as i32).map(|v| if v % 7 == 0 { None } else { Some(v) }), + ); + let str_col = StringArray::from_iter((0..num_rows).map(|v| { + if v % 5 == 0 { + None + } else { + Some(format!("row-{v}")) + } + })); + let bool_col = BooleanArray::from_iter( + (0..num_rows).map(|v| if v % 11 == 0 { None } else { Some(v % 2 == 0) }), + ); + let fsb_col = FixedSizeBinaryArray::try_from_sparse_iter_with_size( + (0..num_rows).map(|v| { + if v % 9 == 0 { + None + } else { + let mut bytes = [0u8; 4]; + bytes.copy_from_slice(&(v as u32).to_le_bytes()); + Some(bytes) + } + }), + fsb_size, + ) + .unwrap(); + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(int_col), + Arc::new(str_col), + Arc::new(bool_col), + Arc::new(fsb_col), + ], + ) + .unwrap(); + + let indices = UInt32Array::from_iter_values((0..num_rows as u32).step_by(2)); + // Prime < indices.len() (150) so the chunking loop in push_batch_with_index + // runs multiple times and the final chunk is a partial batch (150 % 37 = 2). + let target_batch_size = 37; + + let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), target_batch_size); + coalescer + .push_batch_with_indices(batch.clone(), &indices) + .unwrap(); + coalescer.finish_buffered_batch().unwrap(); + + let mut completed = vec![]; + while let Some(b) = coalescer.next_completed_batch() { + completed.push(b); + } + assert!( + completed.len() > 1, + "expected the output to span multiple batches" + ); + for b in &completed[..completed.len() - 1] { + assert_eq!(b.num_rows(), target_batch_size); + } + + let actual = concat_batches(&schema, completed.iter()).unwrap(); + let expected = take_record_batch(&batch, &indices).unwrap(); + assert_eq!(expected, actual); + } + #[test] fn test_push_batch_with_indices_empty_is_noop() { let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])); From 44bbfcfdf9d918f4baca173a47ad4b7c7e383ff7 Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Tue, 20 Oct 2026 10:29:57 -0400 Subject: [PATCH 5/6] add debug test assertions --- arrow-select/src/coalesce/byte_view.rs | 32 ++++++++++++++++++++++++++ arrow-select/src/coalesce/primitive.rs | 16 +++++++++++++ 2 files changed, 48 insertions(+) diff --git a/arrow-select/src/coalesce/byte_view.rs b/arrow-select/src/coalesce/byte_view.rs index f71e5616c9fa..7f781f169c78 100644 --- a/arrow-select/src/coalesce/byte_view.rs +++ b/arrow-select/src/coalesce/byte_view.rs @@ -388,6 +388,10 @@ impl InProgressByteViewArray { 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 { @@ -482,6 +486,10 @@ fn selected_non_inline_bytes( indices: &[u32], max_buffer_size: usize, ) -> Option { + debug_assert!( + indices.iter().all(|&idx| (idx as usize) < src_views.len()), + "selected_non_inline_bytes requires caller-validated in-bounds indices", + ); let mut total = 0usize; for &idx in indices { // SAFETY: caller validated idx < src_views.len(). @@ -650,6 +658,12 @@ impl InProgressArray for InProgressByteViewArray { let index_values = indices.values(); let src_views = typed.views(); let src_buffers = typed.data_buffers(); + debug_assert!( + index_values + .iter() + .all(|&idx| (idx as usize) < src_views.len()), + "copy_rows_by_indices requires caller-validated in-bounds indices", + ); if src_buffers.is_empty() { let current_len = self.views.len(); @@ -711,6 +725,12 @@ impl InProgressArray for InProgressByteViewArray { let current_len = self.views.len(); self.views.reserve(indices.len()); + debug_assert!( + index_values + .iter() + .all(|&idx| (idx as usize) < src_views.len()), + "copy_rows_by_indices requires caller-validated in-bounds indices", + ); // SAFETY: caller validated indices < typed.len() == src_views.len(). unsafe { let base = self.views.spare_capacity_mut().as_mut_ptr().cast::(); @@ -1111,4 +1131,16 @@ mod tests { ); actual.to_data().validate_full().unwrap(); } + + #[test] + #[cfg(debug_assertions)] + #[should_panic(expected = "caller-validated in-bounds indices")] + fn test_copy_rows_by_indices_debug_assert_rejects_oob() { + let source: ArrayRef = Arc::new(BinaryViewArray::from_iter_values( + [b"a".as_slice(), b"b", b"c"].into_iter(), + )); + let indices = UInt32Array::from(vec![0u32, 99]); + let mut in_progress = InProgressByteViewArray::::new(4); + let _ = in_progress.copy_rows_by_indices(&source, &indices); + } } diff --git a/arrow-select/src/coalesce/primitive.rs b/arrow-select/src/coalesce/primitive.rs index 005286aa6ce9..d238f1cfe0b2 100644 --- a/arrow-select/src/coalesce/primitive.rs +++ b/arrow-select/src/coalesce/primitive.rs @@ -227,6 +227,12 @@ impl InProgressArray for InProgressPrimitiveArray 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 // allocated one output slot per index. unsafe { @@ -429,4 +435,14 @@ mod tests { let expected = take(source.as_ref(), &indices, None).unwrap(); assert_eq!(actual.as_ref(), expected.as_ref()); } + + #[test] + #[cfg(debug_assertions)] + #[should_panic(expected = "caller-validated in-bounds indices")] + fn test_copy_rows_by_indices_debug_assert_rejects_oob() { + let source: ArrayRef = Arc::new(Int32Array::from_iter_values(0..4i32)); + let indices = UInt32Array::from(vec![0u32, 99]); + let mut in_progress = InProgressPrimitiveArray::::new(4, DataType::Int32); + let _ = in_progress.copy_rows_by_indices(&source, &indices); + } } From cec88d2c260ad9e78f0616b7f6356d95685100f6 Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Tue, 20 Oct 2026 18:20:14 -0400 Subject: [PATCH 6/6] fix(coalesce): enforce indexed copy safety contracts --- arrow-select/src/coalesce.rs | 70 +++++++++---- arrow-select/src/coalesce/byte_view.rs | 136 +++++++++++++++---------- arrow-select/src/coalesce/primitive.rs | 19 ++-- 3 files changed, 141 insertions(+), 84 deletions(-) diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs index cfd78b6cdfff..46e27b4ca5cc 100644 --- a/arrow-select/src/coalesce.rs +++ b/arrow-select/src/coalesce.rs @@ -22,7 +22,7 @@ //! [`take`]: crate::take::take use crate::filter::{FilterBuilder, FilterPredicate, FilterSelection}; use crate::interleave::interleave_record_batch; -use crate::take::{take, take_bits, take_record_batch}; +use crate::take::{TakeOptions, take, take_arrays_unchecked, take_bits, take_record_batch}; use arrow_array::types::{BinaryViewType, StringViewType, UInt32Type}; use arrow_array::{Array, ArrayRef, BooleanArray, RecordBatch, UInt32Array, downcast_primitive}; use arrow_buffer::{NullBuffer, NullBufferBuilder}; @@ -42,8 +42,13 @@ use fixed_size_binary::InProgressFixedSizeBinaryArray; use generic::GenericInProgressArray; use primitive::InProgressPrimitiveArray; -// Caller must have bounds-checked `indices` against `source_nulls`. -pub(crate) fn append_gathered_nulls( +/// 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, @@ -380,6 +385,9 @@ impl BatchCoalescer { /// Returns `Err(ArrowError::InvalidArgumentError)` if `indices` is a /// non-null `UInt32Array` with any out-of-bounds value. /// + /// Non-null `UInt32Array` indices may use a direct-copy fast path; other + /// supported index arrays use [`take_record_batch`] instead. + /// /// # Example /// ``` /// # use arrow_array::{record_batch, UInt32Array}; @@ -404,30 +412,43 @@ impl BatchCoalescer { batch: RecordBatch, indices: &dyn Array, ) -> Result<(), ArrowError> { - if let Some(u32_arr) = indices.as_any().downcast_ref::() - && u32_arr.null_count() == 0 - { + let u32_indices = indices + .as_any() + .downcast_ref::() + .filter(|array| array.null_count() == 0); + if let Some(u32_arr) = u32_indices { validate_indices_bounds(u32_arr.values(), batch.num_rows())?; } + // Direct indexed copying is a fast path for non-null UInt32 selections + // that are sparse enough for every column (with a stricter ByteView + // threshold). It is only equivalent to `push_batch(take(...))` when the + // materialized batch would not trigger the configured large-batch bypass. 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 indexed_copy_allowed = should_use_indexed_copy(batch.num_rows(), indices.len()) + && !exceeds_byte_view_threshold; 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::() - && u32_arr.null_count() == 0 + if indexed_copy_allowed + && fast_path_allowed + && let Some(u32_arr) = u32_indices { return self.push_batch_with_index(batch, u32_arr); } - let taken_batch = take_record_batch(&batch, indices)?; + + let taken_batch = if let Some(u32_arr) = u32_indices { + // SAFETY: `u32_arr` was bounds-checked against `batch.num_rows()` above, + // which is the length of every column in the record batch. + let columns = unsafe { take_arrays_unchecked(batch.columns(), u32_arr)? }; + // SAFETY: `take_arrays_unchecked` returns one array of `u32_arr.len()` rows + // for every field in `batch.schema()`. + unsafe { RecordBatch::new_unchecked(batch.schema(), columns, u32_arr.len()) } + } else { + take_record_batch(&batch, indices)? + }; self.push_batch(taken_batch) } @@ -468,7 +489,10 @@ impl BatchCoalescer { &sliced }; for (in_progress, column) in self.in_progress_arrays.iter_mut().zip(columns) { - in_progress.copy_rows_by_indices(column, chunk)?; + // SAFETY: `push_batch_with_indices` validated every index against + // `batch.num_rows()`, which equals `column.len()`, and slicing the + // validated indices preserves both bounds and validity. + unsafe { in_progress.copy_rows_by_indices(column, chunk)? }; } self.buffered_rows += take_n; offset += take_n; @@ -992,8 +1016,12 @@ 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`. - fn copy_rows_by_indices( + /// + /// # Safety + /// + /// `indices` must not contain nulls, and every value must be less than + /// `source.len()`. Specialized implementations may use unchecked indexing. + unsafe fn copy_rows_by_indices( &mut self, source: &ArrayRef, indices: &UInt32Array, @@ -1001,7 +1029,11 @@ trait InProgressArray: std::fmt::Debug + Send + Sync { if indices.is_empty() { return Ok(()); } - let taken = take(source.as_ref(), indices, None)?; + let taken = take( + source.as_ref(), + indices, + Some(TakeOptions { check_bounds: true }), + )?; self.set_source(Some(taken)); let result = self.copy_rows(0, indices.len()); self.set_source(None); diff --git a/arrow-select/src/coalesce/byte_view.rs b/arrow-select/src/coalesce/byte_view.rs index 7f781f169c78..ba88db87940a 100644 --- a/arrow-select/src/coalesce/byte_view.rs +++ b/arrow-select/src/coalesce/byte_view.rs @@ -363,24 +363,36 @@ impl InProgressByteViewArray { 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)] - fn append_indexed_views_and_copy_strings( + unsafe 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, - ); + // 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, + ); + } } - fn append_indexed_views_and_copy_strings_with_max_buffer_size( + /// # 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], @@ -481,7 +493,13 @@ impl InProgressByteViewArray { } } -fn selected_non_inline_bytes( +/// Returns the total size of the selected non-inline values, or `None` if a +/// value exceeds `max_buffer_size` or the total overflows `usize`. +/// +/// # Safety +/// +/// Every value in `indices` must be less than `src_views.len()`. +unsafe fn selected_non_inline_bytes( src_views: &[u128], indices: &[u32], max_buffer_size: usize, @@ -638,7 +656,7 @@ impl InProgressArray for InProgressByteViewArray { Ok(()) } - fn copy_rows_by_indices( + unsafe fn copy_rows_by_indices( &mut self, source: &ArrayRef, indices: &UInt32Array, @@ -650,7 +668,9 @@ impl InProgressArray for InProgressByteViewArray { self.ensure_capacity(); if let Some(src_nulls) = typed.nulls() { - append_gathered_nulls(&mut self.nulls, src_nulls, indices); + // 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()); } @@ -686,8 +706,11 @@ impl InProgressArray for InProgressByteViewArray { .is_some_and(|cached| std::ptr::addr_eq(cached.as_ptr(), Arc::as_ptr(source))); if !same_source - && let Some(selected_bytes) = + && let Some(selected_bytes) = unsafe { + // SAFETY: the caller validated every index against `typed.len()`, + // which equals `src_views.len()`. selected_non_inline_bytes(src_views, index_values, i32::MAX as usize) + } { let source_capacity = src_buffers .iter() @@ -696,12 +719,17 @@ impl InProgressArray for InProgressByteViewArray { (source_capacity, selected_bytes.checked_mul(2)) && source_capacity > compaction_threshold { - self.append_indexed_views_and_copy_strings( - src_views, - index_values, - src_buffers, - selected_bytes, - ); + // SAFETY: the caller validated every index against `typed.len()`, + // which equals `src_views.len()`, and both the views and buffers + // come from the same valid ByteView array. + unsafe { + self.append_indexed_views_and_copy_strings( + src_views, + index_values, + src_buffers, + selected_bytes, + ); + } return Ok(()); } } @@ -1029,7 +1057,8 @@ mod tests { let indices = UInt32Array::from(vec![5u32, 0, 2, 4, 1, 1, 3, 0]); let mut in_progress = InProgressByteViewArray::::new(16); - in_progress.copy_rows_by_indices(&source, &indices).unwrap(); + // SAFETY: `indices` contains no nulls and every index is less than `source.len()`. + unsafe { in_progress.copy_rows_by_indices(&source, &indices).unwrap() }; let actual = in_progress.finish().unwrap(); let expected = take(source.as_ref(), &indices, None).unwrap(); @@ -1045,14 +1074,20 @@ mod tests { let mut in_progress = InProgressByteViewArray::::new(64); let dense_indices = UInt32Array::from_iter_values(0..32u32); - in_progress - .copy_rows_by_indices(&source, &dense_indices) - .unwrap(); + // SAFETY: both index arrays contain no nulls and every index is less than + // `source.len()`. + unsafe { + in_progress + .copy_rows_by_indices(&source, &dense_indices) + .unwrap(); + } let after_first = in_progress.size(); - in_progress - .copy_rows_by_indices(&source, &UInt32Array::from(vec![1u32, 2, 4])) - .unwrap(); + unsafe { + in_progress + .copy_rows_by_indices(&source, &UInt32Array::from(vec![1u32, 2, 4])) + .unwrap(); + } let after_second = in_progress.size(); let second_growth = after_second - after_first; @@ -1075,12 +1110,16 @@ mod tests { let indices_b = UInt32Array::from(vec![2u32, 5, 1, 6, 4]); let mut in_progress = InProgressByteViewArray::::new(16); - in_progress - .copy_rows_by_indices(&source_a, &indices_a) - .unwrap(); - in_progress - .copy_rows_by_indices(&source_b, &indices_b) - .unwrap(); + // SAFETY: both index arrays contain no nulls and every index is less than + // its corresponding source length. + unsafe { + in_progress + .copy_rows_by_indices(&source_a, &indices_a) + .unwrap(); + in_progress + .copy_rows_by_indices(&source_b, &indices_b) + .unwrap(); + } let actual = in_progress.finish().unwrap(); let taken_a = take(source_a.as_ref(), &indices_a, None).unwrap(); @@ -1105,18 +1144,23 @@ mod tests { let indices = [0u32, 3, 7]; let max_buffer_size = 80; let typed = source.as_binary_view(); + // SAFETY: every index is less than `typed.views().len()`. let selected_bytes = - selected_non_inline_bytes(typed.views(), &indices, max_buffer_size).unwrap(); + unsafe { selected_non_inline_bytes(typed.views(), &indices, max_buffer_size) }.unwrap(); let mut in_progress = InProgressByteViewArray::::new(16); in_progress.nulls.append_n_non_nulls(indices.len()); - in_progress.append_indexed_views_and_copy_strings_with_max_buffer_size( - typed.views(), - &indices, - typed.data_buffers(), - selected_bytes, - max_buffer_size, - ); + // SAFETY: every index selects a valid view from `typed`, whose views and + // buffers come from the same validated ByteView array. + unsafe { + in_progress.append_indexed_views_and_copy_strings_with_max_buffer_size( + typed.views(), + &indices, + typed.data_buffers(), + selected_bytes, + max_buffer_size, + ); + } let actual = in_progress.finish().unwrap(); let expected = take(source.as_ref(), &UInt32Array::from(indices.to_vec()), None).unwrap(); @@ -1131,16 +1175,4 @@ mod tests { ); actual.to_data().validate_full().unwrap(); } - - #[test] - #[cfg(debug_assertions)] - #[should_panic(expected = "caller-validated in-bounds indices")] - fn test_copy_rows_by_indices_debug_assert_rejects_oob() { - let source: ArrayRef = Arc::new(BinaryViewArray::from_iter_values( - [b"a".as_slice(), b"b", b"c"].into_iter(), - )); - let indices = UInt32Array::from(vec![0u32, 99]); - let mut in_progress = InProgressByteViewArray::::new(4); - let _ = in_progress.copy_rows_by_indices(&source, &indices); - } } diff --git a/arrow-select/src/coalesce/primitive.rs b/arrow-select/src/coalesce/primitive.rs index d238f1cfe0b2..aa0d3534703b 100644 --- a/arrow-select/src/coalesce/primitive.rs +++ b/arrow-select/src/coalesce/primitive.rs @@ -206,7 +206,7 @@ impl InProgressArray for InProgressPrimitiveArray } } - fn copy_rows_by_indices( + unsafe fn copy_rows_by_indices( &mut self, source: &ArrayRef, indices: &UInt32Array, @@ -218,7 +218,9 @@ impl InProgressArray for InProgressPrimitiveArray self.ensure_capacity(); if let Some(src_nulls) = typed.nulls() { - append_gathered_nulls(&mut self.nulls, src_nulls, indices); + // 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()); } @@ -429,20 +431,11 @@ mod tests { ]); let mut in_progress = InProgressPrimitiveArray::::new(32, DataType::Int32); - in_progress.copy_rows_by_indices(&source, &indices).unwrap(); + // SAFETY: `indices` contains no nulls and every index is less than `source.len()`. + unsafe { in_progress.copy_rows_by_indices(&source, &indices).unwrap() }; let actual = in_progress.finish().unwrap(); let expected = take(source.as_ref(), &indices, None).unwrap(); assert_eq!(actual.as_ref(), expected.as_ref()); } - - #[test] - #[cfg(debug_assertions)] - #[should_panic(expected = "caller-validated in-bounds indices")] - fn test_copy_rows_by_indices_debug_assert_rejects_oob() { - let source: ArrayRef = Arc::new(Int32Array::from_iter_values(0..4i32)); - let indices = UInt32Array::from(vec![0u32, 99]); - let mut in_progress = InProgressPrimitiveArray::::new(4, DataType::Int32); - let _ = in_progress.copy_rows_by_indices(&source, &indices); - } }