From 29bd53272e37674a0e817a29b0dec2f52b7e123e Mon Sep 17 00:00:00 2001 From: cl Date: Sat, 18 Apr 2026 16:39:01 +0800 Subject: [PATCH 1/3] arrow-select: optimise coalesced takes for primitive and view arrays ## What - add a direct `push_batch_with_indices` path in `BatchCoalescer` for primitive, `Utf8View`, and `BinaryView` columns when the index array is integer typed and non-null - teach the coalescer internals to copy indexed primitive and view values directly into in-progress output buffers instead of materialising an intermediate taken `RecordBatch` - add dedicated `take:` coalesce benchmarks and new indexed coalescing tests for primitive, `Utf8View`, and `BinaryView` inputs ## How - route supported batches through a direct indices path that chunks the input indices across coalesced output batch boundaries and reuses the existing in-progress array builders - specialise `InProgressPrimitiveArray::copy_indices` to gather values and build the taken null mask directly from the source array - specialise `InProgressByteViewArray::copy_indices` to gather selected views and nulls directly, compute buffer compaction from the selected views, and lazily compute whole-array buffer usage only when the row-copy path needs it - keep unsupported types on the existing `take_record_batch` fallback so the optimisation only applies where the benchmark data shows it is profitable ## Why It Works - the previous `push_batch_with_indices` implementation always paid to allocate and populate a temporary taken `RecordBatch` before coalescing - for primitive and view arrays, the coalescer can write the selected rows straight into its output builders, avoiding that extra batch materialisation and the extra copy it implies - the view-array path remains safe because it preserves the existing reuse-vs-compact behaviour, but bases sparse compaction decisions on the actually selected views rather than the whole source batch ## Tests And Validation - added indexed coalescing tests for mixed primitive, mixed `Utf8View`, mixed `BinaryView`, and `Utf8` fallback behaviour in `arrow-select/src/coalesce.rs` - added `take:` coalesce benchmarks in `arrow/benches/coalesce_kernels.rs` covering primitive, `Utf8View`, `BinaryView`, `Utf8`, and dictionary-backed schemas - validated with: - `cargo test -p arrow-select coalesce --lib` - `cargo clippy -p arrow-select --lib --tests -- -D warnings` - `cargo clippy -p arrow --bench coalesce_kernels --features test_utils -- -D warnings` ## Benchmark Summary - `take: primitive, 8192, nulls: 0, selectivity: 0.01`: `3.5194-3.5796 ms` -> `1.8780-1.9136 ms` - `take: primitive, 8192, nulls: 0.1, selectivity: 0.01`: `5.5208-5.5708 ms` -> `4.0016-4.1647 ms` - `take: primitive, 8192, nulls: 0, selectivity: 0.001`: `23.684-23.813 ms` -> `5.9713-6.0137 ms` - `take: single_utf8view, 8192, nulls: 0, selectivity: 0.01`: `3.0301-3.0830 ms` -> `2.4513-2.4854 ms` - `take: mixed_utf8view (max_string_len=20), 8192, nulls: 0, selectivity: 0.01`: `1.8643-1.8823 ms` -> `1.2706-1.2856 ms` - `take: single_binaryview, 8192, nulls: 0, selectivity: 0.01`: `3.1346-3.2991 ms` -> `2.7578-2.8539 ms` - `take: mixed_binaryview (max_string_len=20), 8192, nulls: 0, selectivity: 0.01`: `1.9634-2.0215 ms` -> `1.4117-1.4383 ms` Signed-off-by: cl --- arrow-select/src/coalesce.rs | 254 +++++++++++++++++++- arrow-select/src/coalesce/byte_view.rs | 160 ++++++++++--- arrow-select/src/coalesce/primitive.rs | 40 +++- arrow/benches/coalesce_kernels.rs | 319 ++++++++++++++++++++++++- 4 files changed, 733 insertions(+), 40 deletions(-) diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs index 8fe88fb8c377..25ab3b928d6d 100644 --- a/arrow-select/src/coalesce.rs +++ b/arrow-select/src/coalesce.rs @@ -23,7 +23,10 @@ use crate::filter::filter_record_batch; use crate::take::take_record_batch; use arrow_array::types::{BinaryViewType, StringViewType}; -use arrow_array::{Array, ArrayRef, BooleanArray, RecordBatch, downcast_primitive}; +use arrow_array::{ + Array, ArrayRef, BooleanArray, RecordBatch, downcast_integer_array, downcast_primitive, +}; +use arrow_buffer::ArrowNativeType; use arrow_schema::{ArrowError, DataType, SchemaRef}; use std::collections::VecDeque; use std::sync::Arc; @@ -272,6 +275,10 @@ impl BatchCoalescer { batch: RecordBatch, indices: &dyn Array, ) -> Result<(), ArrowError> { + if supports_direct_take_coalescing(&batch, indices) { + return self.push_batch_with_indices_direct(batch, indices); + } + // todo: optimize this to avoid materializing (copying the results of take indices to a new batch) let taken_batch = take_record_batch(&batch, indices)?; self.push_batch(taken_batch) @@ -564,6 +571,92 @@ impl BatchCoalescer { pub fn next_completed_batch(&mut self) -> Option { self.completed.pop_front() } + + fn push_batch_with_indices_direct( + &mut self, + batch: RecordBatch, + indices: &dyn Array, + ) -> Result<(), ArrowError> { + if indices.is_empty() { + return Ok(()); + } + + if self.should_materialize_large_take(indices.len()) { + let taken_batch = take_record_batch(&batch, indices)?; + return self.push_batch(taken_batch); + } + + let (_schema, arrays, _num_rows) = batch.into_parts(); + + self.set_sources(arrays)?; + + let result = (|| { + self.copy_index_chunks(indices)?; + Ok(()) + })(); + + self.clear_sources(); + + result + } + + fn should_materialize_large_take(&self, index_count: usize) -> bool { + self.biggest_coalesce_batch_size.is_some_and(|limit| { + index_count > limit && (self.buffered_rows == 0 || self.buffered_rows > limit) + }) + } + + fn set_sources(&mut self, arrays: Vec) -> Result<(), ArrowError> { + if arrays.len() != self.in_progress_arrays.len() { + return Err(ArrowError::InvalidArgumentError(format!( + "Batch has {} columns but BatchCoalescer expects {}", + arrays.len(), + self.in_progress_arrays.len() + ))); + } + + self.in_progress_arrays + .iter_mut() + .zip(arrays) + .for_each(|(in_progress, array)| { + in_progress.set_source(Some(array)); + }); + + Ok(()) + } + + fn clear_sources(&mut self) { + for in_progress in self.in_progress_arrays.iter_mut() { + in_progress.set_source(None); + } + } + + fn copy_index_chunks(&mut self, indices: &dyn Array) -> Result<(), ArrowError> { + let mut offset = 0; + while offset < indices.len() { + let available = self.target_batch_size.saturating_sub(self.buffered_rows); + if available == 0 { + self.finish_buffered_batch()?; + continue; + } + + let chunk_len = available.min(indices.len() - offset); + let chunk = indices.slice(offset, chunk_len); + + for in_progress in self.in_progress_arrays.iter_mut() { + in_progress.copy_indices(chunk.as_ref())?; + } + + self.buffered_rows += chunk_len; + if self.buffered_rows >= self.target_batch_size { + self.finish_buffered_batch()?; + } + + offset += chunk_len; + } + + Ok(()) + } } /// Return a new `InProgressArray` for the given data type @@ -588,6 +681,20 @@ fn create_in_progress_array(data_type: &DataType, batch_size: usize) -> Box bool { + indices.null_count() == 0 + && indices.data_type().is_integer() + && batch + .schema() + .fields() + .iter() + .all(|field| supports_direct_take_data_type(field.data_type())) +} + +fn supports_direct_take_data_type(data_type: &DataType) -> bool { + data_type.is_primitive() || matches!(data_type, DataType::Utf8View | DataType::BinaryView) +} + /// Incrementally builds up arrays /// /// [`GenericInProgressArray`] is the default implementation that buffers @@ -611,6 +718,19 @@ trait InProgressArray: std::fmt::Debug + Send + Sync { /// Return an error if the source array is not set fn copy_rows(&mut self, offset: usize, len: usize) -> Result<(), ArrowError>; + /// Copy rows addressed by `indices` from the current source array. + fn copy_indices(&mut self, indices: &dyn Array) -> Result<(), ArrowError> { + downcast_integer_array!(indices => { + for index in indices.values() { + self.copy_rows(index.as_usize(), 1)?; + } + Ok(()) + }, + d => Err(ArrowError::InvalidArgumentError(format!( + "Internal Error: unsupported index type for coalescing: {d:?}" + )))) + } + /// Finish the currently in-progress array and return it as an `ArrayRef` fn finish(&mut self) -> Result; } @@ -623,8 +743,8 @@ 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, + BinaryViewArray, Float64Array, Int32Array, Int64Array, RecordBatchOptions, StringArray, + StringViewArray, TimestampNanosecondArray, UInt32Array, UInt64Array, make_array, }; use arrow_buffer::BooleanBufferBuilder; use arrow_schema::{DataType, Field, Schema}; @@ -2185,6 +2305,134 @@ mod tests { assert_eq!(expected, actual); } + #[test] + fn test_coalasce_push_batch_with_indices_mixed_primitives() { + let schema = Arc::new(Schema::new(vec![ + Field::new("c0", DataType::Int32, true), + Field::new("c1", DataType::Float64, true), + ])); + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![Some(10), None, Some(30), Some(40)])) as ArrayRef, + Arc::new(Float64Array::from(vec![ + Some(1.5), + Some(2.5), + None, + Some(4.5), + ])) as ArrayRef, + ], + ) + .unwrap(); + + let indices = UInt32Array::from(vec![3, 1, 3, 0]); + let expected = take_record_batch(&batch, &indices).unwrap(); + + let mut coalescer = BatchCoalescer::new(schema, indices.len()); + coalescer.push_batch_with_indices(batch, &indices).unwrap(); + coalescer.finish_buffered_batch().unwrap(); + let actual = coalescer.next_completed_batch().unwrap(); + + assert_eq!(expected, actual); + } + + #[test] + fn test_coalasce_push_batch_with_indices_utf8_view_mixed() { + let schema = Arc::new(Schema::new(vec![ + Field::new("c0", DataType::Int32, true), + Field::new("c1", DataType::Utf8View, true), + ])); + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![Some(10), None, Some(30), Some(40)])) as ArrayRef, + Arc::new(StringViewArray::from_iter([ + Some("alpha"), + None, + Some("a longer utf8 view value"), + Some("delta"), + ])) as ArrayRef, + ], + ) + .unwrap(); + + let indices = UInt32Array::from(vec![3, 1, 3, 0]); + let expected = take_record_batch(&batch, &indices).unwrap(); + + let mut coalescer = BatchCoalescer::new(schema, indices.len()); + coalescer.push_batch_with_indices(batch, &indices).unwrap(); + coalescer.finish_buffered_batch().unwrap(); + let actual = coalescer.next_completed_batch().unwrap(); + + assert_eq!(expected, actual); + } + + #[test] + fn test_coalasce_push_batch_with_indices_utf8_mixed() { + let schema = Arc::new(Schema::new(vec![ + Field::new("c0", DataType::Int32, true), + Field::new("c1", DataType::Utf8, true), + ])); + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![Some(10), None, Some(30), Some(40)])) as ArrayRef, + Arc::new(StringArray::from(vec![ + Some("alpha"), + None, + Some("gamma"), + Some("delta"), + ])) as ArrayRef, + ], + ) + .unwrap(); + + let indices = UInt32Array::from(vec![3, 1, 3, 0]); + let expected = take_record_batch(&batch, &indices).unwrap(); + + let mut coalescer = BatchCoalescer::new(schema, indices.len()); + coalescer.push_batch_with_indices(batch, &indices).unwrap(); + coalescer.finish_buffered_batch().unwrap(); + let actual = coalescer.next_completed_batch().unwrap(); + + assert_eq!(expected, actual); + } + + #[test] + fn test_coalasce_push_batch_with_indices_binary_view_mixed() { + let schema = Arc::new(Schema::new(vec![ + Field::new("c0", DataType::Int32, true), + Field::new("c1", DataType::BinaryView, true), + ])); + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![Some(10), None, Some(30), Some(40)])) as ArrayRef, + Arc::new(BinaryViewArray::from_iter([ + Some(b"alpha" as &[u8]), + None, + Some(b"a longer binary view value" as &[u8]), + Some(b"delta" as &[u8]), + ])) as ArrayRef, + ], + ) + .unwrap(); + + let indices = UInt32Array::from(vec![3, 1, 3, 0]); + let expected = take_record_batch(&batch, &indices).unwrap(); + + let mut coalescer = BatchCoalescer::new(schema, indices.len()); + coalescer.push_batch_with_indices(batch, &indices).unwrap(); + coalescer.finish_buffered_batch().unwrap(); + let actual = coalescer.next_completed_batch().unwrap(); + + assert_eq!(expected, actual); + } + #[test] fn test_push_batch_schema_mismatch_fewer_columns() { // Coalescer expects 0 columns, batch has 1 diff --git a/arrow-select/src/coalesce/byte_view.rs b/arrow-select/src/coalesce/byte_view.rs index 6062cd5e77aa..082467023226 100644 --- a/arrow-select/src/coalesce/byte_view.rs +++ b/arrow-select/src/coalesce/byte_view.rs @@ -18,8 +18,8 @@ use crate::coalesce::InProgressArray; use arrow_array::cast::AsArray; use arrow_array::types::ByteViewType; -use arrow_array::{Array, ArrayRef, GenericByteViewArray}; -use arrow_buffer::{Buffer, NullBufferBuilder}; +use arrow_array::{Array, ArrayRef, GenericByteViewArray, downcast_integer_array}; +use arrow_buffer::{ArrowNativeType, BooleanBuffer, Buffer, NullBuffer, NullBufferBuilder}; use arrow_data::{ByteView, MAX_INLINE_VIEW_LEN}; use arrow_schema::ArrowError; use std::marker::PhantomData; @@ -58,10 +58,35 @@ pub(crate) struct InProgressByteViewArray { struct Source { /// The array to copy form array: ArrayRef, - /// Should the strings from the source array be copied into new buffers? - need_gc: bool, - /// How many bytes were actually used in the source array's buffers? - ideal_buffer_size: usize, + /// Total capacity of the source array's data buffers. + data_buffer_capacity: usize, + /// Lazily computed bytes actually used by the full source array. + whole_array_used_bytes: Option, +} + +impl Source { + fn new(array: ArrayRef) -> Self { + let s = array.as_byte_view::(); + let data_buffer_capacity = s.data_buffers().iter().map(|b| b.capacity()).sum(); + + Self { + array, + data_buffer_capacity, + whole_array_used_bytes: None, + } + } + + fn whole_array_gc_state(&mut self) -> (bool, usize) { + if self.data_buffer_capacity == 0 { + return (false, 0); + } + + let used_bytes = *self + .whole_array_used_bytes + .get_or_insert_with(|| self.array.as_byte_view::().total_buffer_bytes_used()); + let should_compact = used_bytes != 0 && self.data_buffer_capacity > (used_bytes * 2); + (should_compact, used_bytes) + } } // manually implement Debug because ByteViewType doesn't implement Debug @@ -103,6 +128,46 @@ impl InProgressByteViewArray { } } + fn append_indexed_nulls( + &mut self, + indices: &[I], + source_nulls: Option<&NullBuffer>, + ) { + if let Some(source_nulls) = source_nulls.filter(|n| n.null_count() > 0) { + let taken = NullBuffer::new(BooleanBuffer::collect_bool(indices.len(), |idx| { + source_nulls.is_valid(indices[idx].as_usize()) + })); + self.nulls.append_buffer(&taken); + } else { + self.nulls.append_n_non_nulls(indices.len()); + } + } + + fn collect_indexed_views( + &self, + views: &[u128], + indices: &[I], + ) -> Vec { + indices + .iter() + .map(|index| views[index.as_usize()]) + .collect() + } + + fn append_indexed_views(&mut self, views: &[u128], indices: &[I]) { + self.views + .extend(indices.iter().map(|index| views[index.as_usize()])); + } + + fn selected_buffer_bytes(views: &[u128]) -> usize { + views + .iter() + .map(|view| ByteView::from(*view)) + .filter(|view| view.length > MAX_INLINE_VIEW_LEN) + .map(|view| view.length as usize) + .sum() + } + /// Finishes in progress buffer, if any fn finish_current(&mut self) { let Some(next_buffer) = self.current.take() else { @@ -279,40 +344,23 @@ impl InProgressByteViewArray { impl InProgressArray for InProgressByteViewArray { fn set_source(&mut self, source: Option) { - self.source = source.map(|array| { - let s = array.as_byte_view::(); - - let (need_gc, ideal_buffer_size) = if s.data_buffers().is_empty() { - (false, 0) - } else { - let ideal_buffer_size = s.total_buffer_bytes_used(); - // We don't use get_buffer_memory_size here, because gc is for the contents of the - // data buffers, not views and nulls. - let actual_buffer_size = - s.data_buffers().iter().map(|b| b.capacity()).sum::(); - // copying strings is expensive, so only do it if the array is - // sparse (uses at least 2x the memory it needs) - let need_gc = - ideal_buffer_size != 0 && actual_buffer_size > (ideal_buffer_size * 2); - (need_gc, ideal_buffer_size) - }; - - Source { - array, - need_gc, - ideal_buffer_size, - } - }) + self.source = source.map(|array| Source::new::(array)) } fn copy_rows(&mut self, offset: usize, len: usize) -> Result<(), ArrowError> { self.ensure_capacity(); - let source = self.source.take().ok_or_else(|| { + let mut source = self.source.take().ok_or_else(|| { ArrowError::InvalidArgumentError( "Internal Error: InProgressByteViewArray: source not set".to_string(), ) })?; + let (need_gc, ideal_buffer_size) = if source.data_buffer_capacity == 0 { + (false, 0) + } else { + source.whole_array_gc_state::() + }; + // If creating StringViewArray output, ensure input was valid utf8 too let s = source.array.as_byte_view::(); @@ -329,7 +377,7 @@ impl InProgressArray for InProgressByteViewArray { // If there are no data buffers in s (all inlined views), can append the // views/nulls and done - if source.ideal_buffer_size == 0 { + if ideal_buffer_size == 0 { self.views.extend_from_slice(views); self.source = Some(source); return Ok(()); @@ -337,8 +385,8 @@ impl InProgressArray for InProgressByteViewArray { // Copying the strings into a buffer can be time-consuming so // only do it if the array is sparse - if source.need_gc { - self.append_views_and_copy_strings(views, source.ideal_buffer_size, buffers); + if need_gc { + self.append_views_and_copy_strings(views, ideal_buffer_size, buffers); } else { self.append_views_and_update_buffer_index(views, buffers); } @@ -346,6 +394,50 @@ impl InProgressArray for InProgressByteViewArray { Ok(()) } + fn copy_indices(&mut self, indices: &dyn Array) -> Result<(), ArrowError> { + self.ensure_capacity(); + let source = self.source.take().ok_or_else(|| { + ArrowError::InvalidArgumentError( + "Internal Error: InProgressByteViewArray: source not set".to_string(), + ) + })?; + + let s = source.array.as_byte_view::(); + let views = s.views().as_ref(); + + downcast_integer_array!(indices => { + self.append_indexed_nulls(indices.values(), s.nulls()); + + if source.data_buffer_capacity == 0 { + self.append_indexed_views(views, indices.values()); + self.source = Some(source); + return Ok(()); + } + + let selected_views = self.collect_indexed_views(views, indices.values()); + + let selected_buffer_bytes = Self::selected_buffer_bytes(&selected_views); + + if selected_buffer_bytes == 0 { + self.views.extend_from_slice(&selected_views); + } else if source.data_buffer_capacity > (selected_buffer_bytes * 2) { + self.append_views_and_copy_strings( + &selected_views, + selected_buffer_bytes, + s.data_buffers(), + ); + } else { + self.append_views_and_update_buffer_index(&selected_views, s.data_buffers()); + } + + self.source = Some(source); + Ok(()) + }, + d => Err(ArrowError::InvalidArgumentError(format!( + "Internal Error: unsupported index type for byte view coalescing: {d:?}" + )))) + } + fn finish(&mut self) -> Result { self.finish_current(); assert!(self.current.is_none()); diff --git a/arrow-select/src/coalesce/primitive.rs b/arrow-select/src/coalesce/primitive.rs index a7f2fb32ce49..765d9409670d 100644 --- a/arrow-select/src/coalesce/primitive.rs +++ b/arrow-select/src/coalesce/primitive.rs @@ -17,8 +17,8 @@ use crate::coalesce::InProgressArray; use arrow_array::cast::AsArray; -use arrow_array::{Array, ArrayRef, ArrowPrimitiveType, PrimitiveArray}; -use arrow_buffer::{NullBufferBuilder, ScalarBuffer}; +use arrow_array::{Array, ArrayRef, ArrowPrimitiveType, PrimitiveArray, downcast_integer_array}; +use arrow_buffer::{ArrowNativeType, BooleanBuffer, NullBuffer, NullBufferBuilder, ScalarBuffer}; use arrow_schema::{ArrowError, DataType}; use std::fmt::Debug; use std::sync::Arc; @@ -94,6 +94,42 @@ impl InProgressArray for InProgressPrimitiveArray Ok(()) } + fn copy_indices(&mut self, indices: &dyn Array) -> Result<(), ArrowError> { + self.ensure_capacity(); + + let s = self + .source + .as_ref() + .ok_or_else(|| { + ArrowError::InvalidArgumentError( + "Internal Error: InProgressPrimitiveArray: source not set".to_string(), + ) + })? + .as_primitive::(); + + let values = s.values(); + let source_nulls = s.nulls(); + + downcast_integer_array!(indices => { + match source_nulls.filter(|nulls| nulls.null_count() > 0) { + Some(nulls) => { + let nulls = NullBuffer::new(BooleanBuffer::collect_bool(indices.len(), |idx| { + nulls.is_valid(indices.value(idx).as_usize()) + })); + self.nulls.append_buffer(&nulls); + } + None => self.nulls.append_n_non_nulls(indices.len()), + } + + self.current + .extend(indices.values().iter().map(|index| values[index.as_usize()])); + Ok(()) + }, + d => Err(ArrowError::InvalidArgumentError(format!( + "Internal Error: unsupported index type for primitive coalescing: {d:?}" + )))) + } + fn finish(&mut self) -> Result { // take and reset the current values and nulls let values = std::mem::take(&mut self.current); diff --git a/arrow/benches/coalesce_kernels.rs b/arrow/benches/coalesce_kernels.rs index b85c5cc532db..5a79f1ade77f 100644 --- a/arrow/benches/coalesce_kernels.rs +++ b/arrow/benches/coalesce_kernels.rs @@ -25,6 +25,7 @@ use arrow_array::types::{Float64Type, Int32Type, TimestampNanosecondType}; use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit}; use arrow_select::coalesce::BatchCoalescer; use criterion::{Criterion, criterion_group, criterion_main}; +use rand::{Rng, SeedableRng, rngs::StdRng}; /// Benchmarks for generating evently sized output RecordBatches /// from a sequence of filtered source batches @@ -158,7 +159,173 @@ fn add_all_filter_benchmarks(c: &mut Criterion) { } } -criterion_group!(benches, add_all_filter_benchmarks); +fn add_all_take_benchmarks(c: &mut Criterion) { + let batch_size = 8192; + + let primitive_schema = SchemaRef::new(Schema::new(vec![ + Field::new("int32_val", DataType::Int32, true), + Field::new("float_val", DataType::Float64, true), + Field::new( + "timestamp_val", + DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())), + true, + ), + ])); + + let single_schema = SchemaRef::new(Schema::new(vec![Field::new( + "value", + DataType::Utf8View, + true, + )])); + + let single_binaryview_schema = SchemaRef::new(Schema::new(vec![Field::new( + "value", + DataType::BinaryView, + true, + )])); + + let mixed_utf8view_schema = SchemaRef::new(Schema::new(vec![ + Field::new("int32_val", DataType::Int32, true), + Field::new("float_val", DataType::Float64, true), + Field::new("utf8view_val", DataType::Utf8View, true), + ])); + + let mixed_binaryview_schema = SchemaRef::new(Schema::new(vec![ + Field::new("int32_val", DataType::Int32, true), + Field::new("float_val", DataType::Float64, true), + Field::new("binaryview_val", DataType::BinaryView, true), + ])); + + let mixed_utf8_schema = SchemaRef::new(Schema::new(vec![ + Field::new("int32_val", DataType::Int32, true), + Field::new("float_val", DataType::Float64, true), + Field::new("utf8", DataType::Utf8, true), + ])); + + let mixed_dict_schema = SchemaRef::new(Schema::new(vec![ + Field::new( + "string_dict", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + true, + ), + Field::new("float_val1", DataType::Float64, true), + Field::new("float_val2", DataType::Float64, true), + ])); + + for null_density in [0.0, 0.1] { + for selectivity in [0.001, 0.01, 0.1, 0.8] { + TakeBenchmarkBuilder { + c, + name: "primitive", + batch_size, + num_output_batches: 50, + null_density, + selectivity, + max_string_len: 30, + schema: &primitive_schema, + } + .build(); + + TakeBenchmarkBuilder { + c, + name: "single_utf8view", + batch_size, + num_output_batches: 50, + null_density, + selectivity, + max_string_len: 30, + schema: &single_schema, + } + .build(); + + TakeBenchmarkBuilder { + c, + name: "single_binaryview", + batch_size, + num_output_batches: 50, + null_density, + selectivity, + max_string_len: 30, + schema: &single_binaryview_schema, + } + .build(); + + TakeBenchmarkBuilder { + c, + name: "mixed_utf8view (max_string_len=20)", + batch_size, + num_output_batches: 20, + null_density, + selectivity, + max_string_len: 20, + schema: &mixed_utf8view_schema, + } + .build(); + + TakeBenchmarkBuilder { + c, + name: "mixed_utf8view (max_string_len=128)", + batch_size, + num_output_batches: 20, + null_density, + selectivity, + max_string_len: 128, + schema: &mixed_utf8view_schema, + } + .build(); + + TakeBenchmarkBuilder { + c, + name: "mixed_binaryview (max_string_len=20)", + batch_size, + num_output_batches: 20, + null_density, + selectivity, + max_string_len: 20, + schema: &mixed_binaryview_schema, + } + .build(); + + TakeBenchmarkBuilder { + c, + name: "mixed_binaryview (max_string_len=128)", + batch_size, + num_output_batches: 20, + null_density, + selectivity, + max_string_len: 128, + schema: &mixed_binaryview_schema, + } + .build(); + + TakeBenchmarkBuilder { + c, + name: "mixed_utf8", + batch_size, + num_output_batches: 20, + null_density, + selectivity, + max_string_len: 30, + schema: &mixed_utf8_schema, + } + .build(); + + TakeBenchmarkBuilder { + c, + name: "mixed_dict", + batch_size, + num_output_batches: 10, + null_density, + selectivity, + max_string_len: 30, + schema: &mixed_dict_schema, + } + .build(); + } + } +} + +criterion_group!(benches, add_all_filter_benchmarks, add_all_take_benchmarks); criterion_main!(benches); /// Run the filters with a batch_size, null_density, selectivity, and schema @@ -222,6 +389,52 @@ impl FilterBenchmarkBuilder<'_> { } } +struct TakeBenchmarkBuilder<'a> { + c: &'a mut Criterion, + name: &'a str, + batch_size: usize, + num_output_batches: usize, + null_density: f32, + selectivity: f32, + max_string_len: usize, + schema: &'a SchemaRef, +} + +impl TakeBenchmarkBuilder<'_> { + fn build(self) { + let Self { + c, + name, + batch_size, + num_output_batches, + null_density, + selectivity, + max_string_len, + schema, + } = self; + + let indices = IndexStreamBuilder::new() + .with_batch_size(batch_size) + .with_selectivity(selectivity) + .build(); + + let data = DataStreamBuilder::new(Arc::clone(schema)) + .with_batch_size(batch_size) + .with_null_density(null_density) + .with_max_string_len(max_string_len) + .build(); + + let id = format!( + "take: {name}, {batch_size}, nulls: {null_density}, selectivity: {selectivity}" + ); + c.bench_function(&id, |b| { + b.iter(|| { + take_streams(num_output_batches, indices.clone(), data.clone()); + }) + }); + } +} + /// Pull RecordBatches from a data stream and apply a sequence of /// filters from a filter stream until we have a specified number of output /// batches. @@ -247,6 +460,27 @@ fn filter_streams( } } +fn take_streams( + mut num_output_batches: usize, + mut index_stream: IndexStream, + mut data_stream: DataStream, +) { + let schema = data_stream.schema(); + let batch_size = data_stream.batch_size(); + let mut coalescer = BatchCoalescer::new(Arc::clone(schema), batch_size); + + while num_output_batches > 0 { + let indices = index_stream.next_indices(); + let batch = data_stream.next_batch(); + coalescer + .push_batch_with_indices(batch.clone(), indices) + .unwrap(); + if coalescer.next_completed_batch().is_some() { + num_output_batches -= 1; + } + } +} + /// Stream of filters to apply to a sequence of input RecordBatches /// /// This pre-computes a sequence of filters and then repeats it forever. @@ -257,6 +491,25 @@ struct FilterStream { batches: Arc<[BooleanArray]>, } +#[derive(Debug, Clone)] +struct IndexStream { + index: usize, + batches: Arc<[UInt32Array]>, +} + +impl IndexStream { + fn next_indices(&mut self) -> &UInt32Array { + let current_index = self.index; + self.index += 1; + if self.index >= self.batches.len() { + self.index = 0; + } + self.batches + .get(current_index) + .expect("No more index batches available") + } +} + impl FilterStream { pub fn next_filter(&mut self) -> &BooleanArray { let current_index = self.index; @@ -324,6 +577,59 @@ impl FilterStreamBuilder { } } +#[derive(Debug)] +struct IndexStreamBuilder { + batch_size: usize, + num_batches: usize, + selectivity: f32, +} + +impl IndexStreamBuilder { + fn new() -> Self { + Self { + batch_size: 8192, + num_batches: 11, + selectivity: 0.5, + } + } + + fn with_batch_size(mut self, batch_size: usize) -> Self { + self.batch_size = batch_size; + self + } + + fn with_selectivity(mut self, selectivity: f32) -> Self { + assert!((0.0..=1.0).contains(&selectivity)); + self.selectivity = selectivity; + self + } + + fn build(self) -> IndexStream { + let Self { + batch_size, + num_batches, + selectivity, + } = self; + + let output_len = ((batch_size as f32) * selectivity).round().max(1.0) as usize; + let batches = (0..num_batches) + .map(|seed| create_index_array(batch_size, output_len, seed as u64)) + .collect::>(); + + IndexStream { + index: 0, + batches: Arc::from(batches), + } + } +} + +fn create_index_array(input_len: usize, output_len: usize, seed: u64) -> UInt32Array { + let mut rng = StdRng::seed_from_u64(seed); + UInt32Array::from_iter_values( + (0..output_len).map(|_| rng.random_range(0..u32::try_from(input_len).unwrap())), + ) +} + #[derive(Debug, Clone)] struct DataStream { schema: SchemaRef, @@ -455,6 +761,17 @@ impl DataStreamBuilder { self.max_string_len, )) // TODO seed } + DataType::BinaryView => Arc::new(BinaryViewArray::from_iter( + create_binary_array_with_len_range_and_prefix_and_seed::( + self.batch_size, + self.null_density, + 0, + self.max_string_len, + b"", + seed, + ) + .iter(), + )), DataType::Dictionary(key_type, value_type) if key_type.as_ref() == &DataType::Int32 && value_type.as_ref() == &DataType::Utf8 => From 0bdc485bebb6169cfd876ef7d9cd54b433c9a3f9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=94=A1=E7=95=A5?= Date: Sun, 19 Apr 2026 01:17:32 +0800 Subject: [PATCH 2/3] arrow-select: chunk oversized materialized takes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 蔡略 --- arrow-select/src/coalesce.rs | 107 ++++++++-- arrow-select/src/coalesce/generic.rs | 9 +- arrow/benches/coalesce_kernels.rs | 309 +++++++++++++++++---------- 3 files changed, 287 insertions(+), 138 deletions(-) diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs index 25ab3b928d6d..912ca19c838d 100644 --- a/arrow-select/src/coalesce.rs +++ b/arrow-select/src/coalesce.rs @@ -23,10 +23,7 @@ use crate::filter::filter_record_batch; use crate::take::take_record_batch; use arrow_array::types::{BinaryViewType, StringViewType}; -use arrow_array::{ - Array, ArrayRef, BooleanArray, RecordBatch, downcast_integer_array, downcast_primitive, -}; -use arrow_buffer::ArrowNativeType; +use arrow_array::{Array, ArrayRef, BooleanArray, RecordBatch, downcast_primitive}; use arrow_schema::{ArrowError, DataType, SchemaRef}; use std::collections::VecDeque; use std::sync::Arc; @@ -275,13 +272,13 @@ impl BatchCoalescer { batch: RecordBatch, indices: &dyn Array, ) -> Result<(), ArrowError> { + // Only arrays with specialized `copy_indices` implementations use the + // direct path. All other schemas stay on the existing take fallback. if supports_direct_take_coalescing(&batch, indices) { return self.push_batch_with_indices_direct(batch, indices); } - // todo: optimize this to avoid materializing (copying the results of take indices to a new batch) - let taken_batch = take_record_batch(&batch, indices)?; - self.push_batch(taken_batch) + self.push_batch_with_indices_materialized(batch, indices) } /// Push all the rows from `batch` into the Coalescer @@ -582,8 +579,7 @@ impl BatchCoalescer { } if self.should_materialize_large_take(indices.len()) { - let taken_batch = take_record_batch(&batch, indices)?; - return self.push_batch(taken_batch); + return self.push_batch_with_indices_materialized(batch, indices); } let (_schema, arrays, _num_rows) = batch.into_parts(); @@ -600,6 +596,49 @@ impl BatchCoalescer { result } + fn push_batch_with_indices_materialized( + &mut self, + batch: RecordBatch, + indices: &dyn Array, + ) -> Result<(), ArrowError> { + if indices.is_empty() { + return Ok(()); + } + + let Some(limit) = self.biggest_coalesce_batch_size.filter(|limit| *limit > 0) else { + let taken_batch = take_record_batch(&batch, indices)?; + return self.push_batch(taken_batch); + }; + + // Match the large-batch path by flushing oversized buffered state + // before materializing additional index chunks. + if self.buffered_rows > limit { + self.finish_buffered_batch()?; + } + + let mut offset = 0; + while offset < indices.len() { + let remaining_indices = indices.len() - offset; + let chunk_len = self.next_materialized_take_chunk_len(remaining_indices, limit); + let chunk = indices.slice(offset, chunk_len); + let taken_batch = take_record_batch(&batch, chunk.as_ref())?; + self.push_batch(taken_batch)?; + offset += chunk_len; + } + + Ok(()) + } + + /// Fill the current buffered batch first, then keep subsequent materialized + /// takes under `limit` rows. + fn next_materialized_take_chunk_len(&self, remaining_indices: usize, limit: usize) -> usize { + if self.buffered_rows > 0 && self.buffered_rows < limit { + return (limit - self.buffered_rows).min(remaining_indices); + } + + limit.min(remaining_indices) + } + fn should_materialize_large_take(&self, index_count: usize) -> bool { self.biggest_coalesce_batch_size.is_some_and(|limit| { index_count > limit && (self.buffered_rows == 0 || self.buffered_rows > limit) @@ -719,17 +758,11 @@ trait InProgressArray: std::fmt::Debug + Send + Sync { fn copy_rows(&mut self, offset: usize, len: usize) -> Result<(), ArrowError>; /// Copy rows addressed by `indices` from the current source array. - fn copy_indices(&mut self, indices: &dyn Array) -> Result<(), ArrowError> { - downcast_integer_array!(indices => { - for index in indices.values() { - self.copy_rows(index.as_usize(), 1)?; - } - Ok(()) - }, - d => Err(ArrowError::InvalidArgumentError(format!( - "Internal Error: unsupported index type for coalescing: {d:?}" - )))) - } + /// + /// This is only used by the direct indexed coalescing fast path for + /// primitive and byte-view arrays. Other data types stay on the + /// `take_record_batch` fallback. + fn copy_indices(&mut self, indices: &dyn Array) -> Result<(), ArrowError>; /// Finish the currently in-progress array and return it as an `ArrayRef` fn finish(&mut self) -> Result; @@ -2433,6 +2466,40 @@ mod tests { assert_eq!(expected, actual); } + #[test] + fn test_coalasce_push_batch_with_indices_respects_biggest_coalesce_batch_size() { + assert_push_batch_with_indices_respects_biggest_coalesce_batch_size(create_test_batch( + 1200, + )); + } + + #[test] + fn test_coalasce_push_batch_with_indices_utf8_respects_biggest_coalesce_batch_size() { + assert_push_batch_with_indices_respects_biggest_coalesce_batch_size(utf8_batch(0..1200)); + } + + fn assert_push_batch_with_indices_respects_biggest_coalesce_batch_size(batch: RecordBatch) { + let schema = batch.schema(); + let indices = UInt32Array::from_iter_values(0..1200_u32); + let expected = take_record_batch(&batch, &indices).unwrap(); + + let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), 1000) + .with_biggest_coalesce_batch_size(Some(500)); + coalescer.push_batch_with_indices(batch, &indices).unwrap(); + + let first = coalescer.next_completed_batch().unwrap(); + assert_eq!(first.num_rows(), 1000); + assert!(!coalescer.has_completed_batch()); + + coalescer.finish_buffered_batch().unwrap(); + let second = coalescer.next_completed_batch().unwrap(); + assert_eq!(second.num_rows(), 200); + assert!(coalescer.next_completed_batch().is_none()); + + let actual = concat_batches(&schema, &[first, second]).unwrap(); + assert_eq!(expected, actual); + } + #[test] fn test_push_batch_schema_mismatch_fewer_columns() { // Coalescer expects 0 columns, batch has 1 diff --git a/arrow-select/src/coalesce/generic.rs b/arrow-select/src/coalesce/generic.rs index 1ea57dff929c..0e1d91087884 100644 --- a/arrow-select/src/coalesce/generic.rs +++ b/arrow-select/src/coalesce/generic.rs @@ -17,7 +17,7 @@ use super::InProgressArray; use crate::concat::concat; -use arrow_array::ArrayRef; +use arrow_array::{Array, ArrayRef}; use arrow_schema::ArrowError; /// Generic implementation for [`InProgressArray`] that works with any type of @@ -60,6 +60,13 @@ impl InProgressArray for GenericInProgressArray { Ok(()) } + fn copy_indices(&mut self, _indices: &dyn Array) -> Result<(), ArrowError> { + Err(ArrowError::InvalidArgumentError( + "Internal Error: indexed coalescing is only supported for primitive and byte view arrays" + .to_string(), + )) + } + fn finish(&mut self) -> Result { // Concatenate all buffered arrays into a single array, which uses 2x // peak memory diff --git a/arrow/benches/coalesce_kernels.rs b/arrow/benches/coalesce_kernels.rs index 5a79f1ade77f..0816d1a2e8bb 100644 --- a/arrow/benches/coalesce_kernels.rs +++ b/arrow/benches/coalesce_kernels.rs @@ -214,114 +214,96 @@ fn add_all_take_benchmarks(c: &mut Criterion) { for null_density in [0.0, 0.1] { for selectivity in [0.001, 0.01, 0.1, 0.8] { - TakeBenchmarkBuilder { - c, - name: "primitive", - batch_size, - num_output_batches: 50, - null_density, - selectivity, - max_string_len: 30, - schema: &primitive_schema, - } - .build(); - - TakeBenchmarkBuilder { - c, - name: "single_utf8view", - batch_size, - num_output_batches: 50, - null_density, - selectivity, - max_string_len: 30, - schema: &single_schema, - } - .build(); - - TakeBenchmarkBuilder { - c, - name: "single_binaryview", - batch_size, - num_output_batches: 50, - null_density, - selectivity, - max_string_len: 30, - schema: &single_binaryview_schema, - } - .build(); - - TakeBenchmarkBuilder { - c, - name: "mixed_utf8view (max_string_len=20)", - batch_size, - num_output_batches: 20, - null_density, - selectivity, - max_string_len: 20, - schema: &mixed_utf8view_schema, - } - .build(); - - TakeBenchmarkBuilder { - c, - name: "mixed_utf8view (max_string_len=128)", - batch_size, - num_output_batches: 20, - null_density, - selectivity, - max_string_len: 128, - schema: &mixed_utf8view_schema, - } - .build(); - - TakeBenchmarkBuilder { - c, - name: "mixed_binaryview (max_string_len=20)", - batch_size, - num_output_batches: 20, - null_density, - selectivity, - max_string_len: 20, - schema: &mixed_binaryview_schema, - } - .build(); - - TakeBenchmarkBuilder { - c, - name: "mixed_binaryview (max_string_len=128)", - batch_size, - num_output_batches: 20, - null_density, - selectivity, - max_string_len: 128, - schema: &mixed_binaryview_schema, - } - .build(); - - TakeBenchmarkBuilder { - c, - name: "mixed_utf8", - batch_size, - num_output_batches: 20, - null_density, - selectivity, - max_string_len: 30, - schema: &mixed_utf8_schema, + for scenario in [ + TakeBenchmarkScenario { + name: "primitive", + num_output_batches: 50, + max_string_len: 30, + schema: &primitive_schema, + }, + TakeBenchmarkScenario { + name: "single_utf8view", + num_output_batches: 50, + max_string_len: 30, + schema: &single_schema, + }, + TakeBenchmarkScenario { + name: "single_binaryview", + num_output_batches: 50, + max_string_len: 30, + schema: &single_binaryview_schema, + }, + TakeBenchmarkScenario { + name: "mixed_utf8view (max_string_len=20)", + num_output_batches: 20, + max_string_len: 20, + schema: &mixed_utf8view_schema, + }, + TakeBenchmarkScenario { + name: "mixed_utf8view (max_string_len=128)", + num_output_batches: 20, + max_string_len: 128, + schema: &mixed_utf8view_schema, + }, + TakeBenchmarkScenario { + name: "mixed_binaryview (max_string_len=20)", + num_output_batches: 20, + max_string_len: 20, + schema: &mixed_binaryview_schema, + }, + TakeBenchmarkScenario { + name: "mixed_binaryview (max_string_len=128)", + num_output_batches: 20, + max_string_len: 128, + schema: &mixed_binaryview_schema, + }, + TakeBenchmarkScenario { + name: "mixed_utf8", + num_output_batches: 20, + max_string_len: 30, + schema: &mixed_utf8_schema, + }, + TakeBenchmarkScenario { + name: "mixed_dict", + num_output_batches: 10, + max_string_len: 30, + schema: &mixed_dict_schema, + }, + ] { + TakeBenchmarkBuilder::from_scenario( + c, + batch_size, + null_density, + selectivity, + scenario, + ) + .build(); } - .build(); + } + } - TakeBenchmarkBuilder { - c, - name: "mixed_dict", - batch_size, - num_output_batches: 10, - null_density, - selectivity, + // Repeated indices make the taken batch much larger than the source batch, + // which exercises the materialized fallback path for unsupported schemas. + for (name, schema) in [ + ("primitive extra_large_repeat", &primitive_schema), + ("mixed_utf8 extra_large_repeat", &mixed_utf8_schema), + ] { + TakeBenchmarkBuilder::from_scenario( + c, + batch_size, + 0.0, + 1.0, + TakeBenchmarkScenario { + name, + num_output_batches: 64, max_string_len: 30, - schema: &mixed_dict_schema, - } - .build(); - } + schema, + }, + ) + .with_biggest_coalesce_batch_size(1024) + .with_index_output_len(131_072) + .with_drain_all_completed_batches() + .build(); } } @@ -398,9 +380,64 @@ struct TakeBenchmarkBuilder<'a> { selectivity: f32, max_string_len: usize, schema: &'a SchemaRef, + biggest_coalesce_batch_size: Option, + index_output_len: Option, + drain_all_completed_batches: bool, } -impl TakeBenchmarkBuilder<'_> { +#[derive(Clone, Copy)] +struct TakeBenchmarkScenario<'a> { + name: &'a str, + num_output_batches: usize, + max_string_len: usize, + schema: &'a SchemaRef, +} + +impl<'a> TakeBenchmarkBuilder<'a> { + fn from_scenario( + c: &'a mut Criterion, + batch_size: usize, + null_density: f32, + selectivity: f32, + scenario: TakeBenchmarkScenario<'a>, + ) -> TakeBenchmarkBuilder<'a> { + let TakeBenchmarkScenario { + name, + num_output_batches, + max_string_len, + schema, + } = scenario; + + TakeBenchmarkBuilder { + c, + name, + batch_size, + num_output_batches, + null_density, + selectivity, + max_string_len, + schema, + biggest_coalesce_batch_size: None, + index_output_len: None, + drain_all_completed_batches: false, + } + } + + fn with_biggest_coalesce_batch_size(mut self, limit: usize) -> Self { + self.biggest_coalesce_batch_size = Some(limit); + self + } + + fn with_index_output_len(mut self, output_len: usize) -> Self { + self.index_output_len = Some(output_len); + self + } + + fn with_drain_all_completed_batches(mut self) -> Self { + self.drain_all_completed_batches = true; + self + } + fn build(self) { let Self { c, @@ -411,12 +448,24 @@ impl TakeBenchmarkBuilder<'_> { selectivity, max_string_len, schema, + biggest_coalesce_batch_size, + index_output_len, + drain_all_completed_batches, } = self; - let indices = IndexStreamBuilder::new() - .with_batch_size(batch_size) - .with_selectivity(selectivity) - .build(); + let output_len = index_output_len + .unwrap_or_else(|| ((batch_size as f32) * selectivity).round().max(1.0) as usize); + + let indices = match index_output_len { + Some(_) => IndexStreamBuilder::new() + .with_batch_size(batch_size) + .with_output_len(output_len) + .build(), + None => IndexStreamBuilder::new() + .with_batch_size(batch_size) + .with_selectivity(selectivity) + .build(), + }; let data = DataStreamBuilder::new(Arc::clone(schema)) .with_batch_size(batch_size) @@ -424,12 +473,22 @@ impl TakeBenchmarkBuilder<'_> { .with_max_string_len(max_string_len) .build(); - let id = format!( - "take: {name}, {batch_size}, nulls: {null_density}, selectivity: {selectivity}" - ); + let id = if index_output_len.is_some() || biggest_coalesce_batch_size.is_some() { + format!( + "take: {name}, input: {batch_size}, output: {output_len}, nulls: {null_density}, biggest: {biggest_coalesce_batch_size:?}" + ) + } else { + format!("take: {name}, {batch_size}, nulls: {null_density}, selectivity: {selectivity}") + }; c.bench_function(&id, |b| { b.iter(|| { - take_streams(num_output_batches, indices.clone(), data.clone()); + take_streams( + num_output_batches, + indices.clone(), + data.clone(), + biggest_coalesce_batch_size, + drain_all_completed_batches, + ); }) }); } @@ -464,10 +523,13 @@ fn take_streams( mut num_output_batches: usize, mut index_stream: IndexStream, mut data_stream: DataStream, + biggest_coalesce_batch_size: Option, + drain_all_completed_batches: bool, ) { let schema = data_stream.schema(); let batch_size = data_stream.batch_size(); - let mut coalescer = BatchCoalescer::new(Arc::clone(schema), batch_size); + let mut coalescer = BatchCoalescer::new(Arc::clone(schema), batch_size) + .with_biggest_coalesce_batch_size(biggest_coalesce_batch_size); while num_output_batches > 0 { let indices = index_stream.next_indices(); @@ -475,7 +537,11 @@ fn take_streams( coalescer .push_batch_with_indices(batch.clone(), indices) .unwrap(); - if coalescer.next_completed_batch().is_some() { + if drain_all_completed_batches { + while num_output_batches > 0 && coalescer.next_completed_batch().is_some() { + num_output_batches -= 1; + } + } else if coalescer.next_completed_batch().is_some() { num_output_batches -= 1; } } @@ -581,6 +647,7 @@ impl FilterStreamBuilder { struct IndexStreamBuilder { batch_size: usize, num_batches: usize, + output_len: Option, selectivity: f32, } @@ -589,6 +656,7 @@ impl IndexStreamBuilder { Self { batch_size: 8192, num_batches: 11, + output_len: None, selectivity: 0.5, } } @@ -604,14 +672,21 @@ impl IndexStreamBuilder { self } + fn with_output_len(mut self, output_len: usize) -> Self { + self.output_len = Some(output_len.max(1)); + self + } + fn build(self) -> IndexStream { let Self { batch_size, num_batches, + output_len, selectivity, } = self; - let output_len = ((batch_size as f32) * selectivity).round().max(1.0) as usize; + let output_len = output_len + .unwrap_or_else(|| ((batch_size as f32) * selectivity).round().max(1.0) as usize); let batches = (0..num_batches) .map(|seed| create_index_array(batch_size, output_len, seed as u64)) .collect::>(); From 67ca60c89404cef2f26eaf797a7e476aa7017519 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=94=A1=E7=95=A5?= Date: Sun, 19 Apr 2026 01:51:45 +0800 Subject: [PATCH 3/3] arrow-select: shorten coalesce take test names --- arrow-select/src/coalesce.rs | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs index 912ca19c838d..c4ef02fe5790 100644 --- a/arrow-select/src/coalesce.rs +++ b/arrow-select/src/coalesce.rs @@ -2467,18 +2467,16 @@ mod tests { } #[test] - fn test_coalasce_push_batch_with_indices_respects_biggest_coalesce_batch_size() { - assert_push_batch_with_indices_respects_biggest_coalesce_batch_size(create_test_batch( - 1200, - )); + fn test_push_batch_with_indices_chunks_large_take() { + assert_chunked_push_batch_with_indices(create_test_batch(1200)); } #[test] - fn test_coalasce_push_batch_with_indices_utf8_respects_biggest_coalesce_batch_size() { - assert_push_batch_with_indices_respects_biggest_coalesce_batch_size(utf8_batch(0..1200)); + fn test_push_batch_with_indices_chunks_large_take_utf8() { + assert_chunked_push_batch_with_indices(utf8_batch(0..1200)); } - fn assert_push_batch_with_indices_respects_biggest_coalesce_batch_size(batch: RecordBatch) { + fn assert_chunked_push_batch_with_indices(batch: RecordBatch) { let schema = batch.schema(); let indices = UInt32Array::from_iter_values(0..1200_u32); let expected = take_record_batch(&batch, &indices).unwrap();