diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs index 9fe874ea8bbf..1064b753aef7 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::{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}; use arrow_schema::{ArrowError, DataType, SchemaRef}; use std::collections::VecDeque; use std::sync::Arc; @@ -41,6 +42,54 @@ use fixed_size_binary::InProgressFixedSizeBinaryArray; use generic::GenericInProgressArray; 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::(source_nulls.inner(), indices); + 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 +105,20 @@ 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; + +// 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 @@ -159,6 +222,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 @@ -180,6 +245,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() @@ -191,6 +260,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, @@ -312,15 +382,21 @@ 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. + /// + /// 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, 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); @@ -336,11 +412,97 @@ 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) - let taken_batch = take_record_batch(&batch, indices)?; + 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()); + 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 indexed_copy_allowed + && fast_path_allowed + && let Some(u32_arr) = u32_indices + { + return self.push_batch_with_index(batch, u32_arr); + } + + 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) } + fn push_batch_with_index( + &mut self, + batch: RecordBatch, + indices: &UInt32Array, + ) -> 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() + ))); + } + 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 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) { + // 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; + 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 @@ -861,6 +1023,31 @@ trait InProgressArray: std::fmt::Debug + Send + Sync { } } + /// Copy rows at `indices` from `source` into the in-progress array. + /// + /// # 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, + ) -> Result<(), ArrowError> { + if indices.is_empty() { + return Ok(()); + } + 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); + result + } + /// Finish the currently in-progress array and return it as an `ArrayRef` fn finish(&mut self) -> Result; @@ -873,13 +1060,14 @@ 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, UInt8Type}; use arrow_array::{ - BinaryViewArray, DictionaryArray, Int32Array, Int64Array, RecordBatchOptions, StringArray, - StringViewArray, TimestampNanosecondArray, UInt8Array, UInt32Array, UInt64Array, - make_array, record_batch, + BinaryViewArray, DictionaryArray, FixedSizeBinaryArray, Int32Array, Int64Array, + RecordBatchOptions, StringArray, StringViewArray, TimestampNanosecondArray, UInt8Array, + UInt32Array, UInt64Array, make_array, record_batch, }; use arrow_buffer::BooleanBufferBuilder; use arrow_schema::{DataType, Field, Schema}; @@ -974,6 +1162,19 @@ 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_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 @@ -3043,4 +3244,296 @@ 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_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)])); + 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_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( + "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 085bf1f1e78c..a5d39ad56177 100644 --- a/arrow-select/src/coalesce/byte_view.rs +++ b/arrow-select/src/coalesce/byte_view.rs @@ -15,16 +15,16 @@ // 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; -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; 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,166 @@ impl InProgressByteViewArray { 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::(); + 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"); + 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"); + } + // 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); + } +} + +/// 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, +) -> 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(). + 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 +529,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::(); @@ -495,6 +664,137 @@ impl InProgressArray for InProgressByteViewArray { Ok(()) } + unsafe fn copy_rows_by_indices( + &mut self, + source: &ArrayRef, + indices: &UInt32Array, + ) -> Result<(), ArrowError> { + if indices.is_empty() { + return Ok(()); + } + let typed = source.as_byte_view::(); + 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 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(); + 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 index_values.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) = 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() + .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 + { + // 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(()); + } + } + + 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()); + 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::(); + let src = src_views.as_ptr(); + if starting_buffer == 0 { + 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 index_values.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()); @@ -506,6 +806,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 @@ -570,9 +872,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() { @@ -745,4 +1049,138 @@ 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 = UInt32Array::from(vec![5u32, 0, 2, 4, 1, 1, 3, 0]); + + let mut in_progress = InProgressByteViewArray::::new(16); + // 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()); + 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 = UInt32Array::from_iter_values(0..32u32); + // 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(); + + 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; + 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 = 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); + // 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(); + 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()); + + 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(); + // SAFETY: every index is less than `typed.views().len()`. + let selected_bytes = + 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()); + // 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(); + 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..aa0d3534703b 100644 --- a/arrow-select/src/coalesce/primitive.rs +++ b/arrow-select/src/coalesce/primitive.rs @@ -15,10 +15,10 @@ // 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}; +use arrow_array::{Array, ArrayRef, ArrowPrimitiveType, PrimitiveArray, UInt32Array}; use arrow_buffer::{NullBuffer, NullBufferBuilder, ScalarBuffer}; use arrow_schema::{ArrowError, DataType}; use std::fmt::Debug; @@ -206,6 +206,52 @@ impl 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::(); + 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 + // 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 index_values.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 +277,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 +420,22 @@ 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 = 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); + // 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()); + } } diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index 2b01a4b1f528..ef3a6f36e420 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 {