diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index ecff2edbbec16..2cd7bd61b648b 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs @@ -2839,10 +2839,11 @@ async fn test_sort_with_streaming_table() -> Result<()> { let sql = "SELECT a FROM test_table GROUP BY a ORDER BY a"; let results = ctx.sql(sql).await?.collect().await?; - assert_eq!(results.len(), 1); - assert_eq!(results[0].num_columns(), 1); + // the merge may split the output into several batches, so concat first + let results = arrow::compute::concat_batches(&results[0].schema(), &results)?; + assert_eq!(results.num_columns(), 1); let expected = create_array!(Int32, vec![1, 2, 3]) as ArrayRef; - assert_eq!(results[0].column(0), &expected); + assert_eq!(results.column(0), &expected); Ok(()) } diff --git a/datafusion/ffi/tests/ffi_udaf.rs b/datafusion/ffi/tests/ffi_udaf.rs index 7df3404d7421b..441f8ddc4f88c 100644 --- a/datafusion/ffi/tests/ffi_udaf.rs +++ b/datafusion/ffi/tests/ffi_udaf.rs @@ -56,13 +56,16 @@ mod tests { .sort_by(vec![col("a")])?; let result = df.collect().await?; + // the merge may split the output into several batches, so concat first + let result = + arrow::compute::concat_batches(&result[0].schema(), &result).unwrap(); let expected = record_batch!( ("a", Int32, vec![1, 2, 4]), ("sum_b", Float64, vec![1.0, 4.0, 16.0]) )?; - assert_eq!(result[0], expected); + assert_eq!(result, expected); Ok(()) } @@ -105,7 +108,10 @@ mod tests { .sort_by(vec![col("a")])?; let result = df.collect().await?; - let result = result[0].column_by_name("stddev_b").unwrap(); + // the merge may split the output into several batches, so concat first + let result = + arrow::compute::concat_batches(&result[0].schema(), &result).unwrap(); + let result = result.column_by_name("stddev_b").unwrap(); let result = result .as_any() .downcast_ref::() @@ -162,13 +168,16 @@ mod tests { df.clone().show().await?; let result = df.collect().await?; + // the merge may split the output into several batches, so concat first + let result = + arrow::compute::concat_batches(&result[0].schema(), &result).unwrap(); let expected = record_batch!( ("a", Int32, vec![1, 2, 4]), ("sum_b", Float64, vec![2.0, 8.0, 32.0]) )?; - assert_eq!(result[0], expected); + assert_eq!(result, expected); Ok(()) } diff --git a/datafusion/physical-plan/benches/sort_preserving_merge.rs b/datafusion/physical-plan/benches/sort_preserving_merge.rs index 76ebf230a30e0..76e39432b1a4a 100644 --- a/datafusion/physical-plan/benches/sort_preserving_merge.rs +++ b/datafusion/physical-plan/benches/sort_preserving_merge.rs @@ -21,15 +21,28 @@ use arrow::{ }; use arrow_schema::{SchemaRef, SortOptions}; use criterion::{BatchSize, Criterion, criterion_group, criterion_main}; +use datafusion_execution::SendableRecordBatchStream; use datafusion_execution::TaskContext; use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr, expressions::col}; use datafusion_physical_plan::test::TestMemoryExec; use datafusion_physical_plan::{ - collect, sorts::sort_preserving_merge::SortPreservingMergeExec, + collect, execute_stream, sorts::sort_preserving_merge::SortPreservingMergeExec, }; +use futures::StreamExt; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; +use std::hint::black_box; use std::sync::Arc; +/// Consume the stream batch by batch, dropping each batch as it arrives +/// instead of holding the whole result in memory like `collect` would +async fn drain(mut stream: SendableRecordBatchStream) { + while let Some(batch) = stream.next().await { + black_box(batch.unwrap()); + } +} + const BENCH_ROWS: usize = 1_000_000; // 1 million rows fn get_large_string(idx: usize) -> String { @@ -193,5 +206,161 @@ fn bench_merge_sorted_preserving(c: &mut Criterion) { } } -criterion_group!(benches, bench_merge_sorted_preserving); +// --------------------------------------------------------------------------- +// Benchmarks across data orderings (sorted / nearly sorted / reverse / +// unsorted), sort key types (u64 / string / complex) and payload widths +// (5 / 20 / 100 columns). +// --------------------------------------------------------------------------- + +const NUM_PARTITIONS: usize = 4; +const ROWS_PER_PARTITION: usize = 100_000; +const BATCH_SIZE: usize = 8192; + +const ORDERINGS: [&str; 4] = ["sorted", "nearly_sorted", "reverse", "unsorted"]; +const KEY_TYPES: [&str; 3] = ["u64", "string", "complex"]; +const PAYLOAD_WIDTHS: [usize; 3] = [5, 20, 100]; + +/// Generate the keys in their "arrival" order, before partitioning +fn generate_keys(ordering: &str) -> Vec { + let n = (NUM_PARTITIONS * ROWS_PER_PARTITION) as u64; + let mut rng = StdRng::seed_from_u64(42); + match ordering { + "sorted" => (0..n).collect(), + "reverse" => (0..n).rev().collect(), + // Sorted except for ~1% of items misplaced to random positions + "nearly_sorted" => { + let mut keys: Vec = (0..n).collect(); + for _ in 0..(n / 100) { + let a = rng.random_range(0..n as usize); + let b = rng.random_range(0..n as usize); + keys.swap(a, b); + } + keys + } + "unsorted" => (0..n).map(|_| rng.random_range(0..n)).collect(), + _ => unreachable!(), + } +} + +/// Distribute the arrival sequence round-robin (a batch at a time) over the +/// partitions, then sort each partition's keys, as SortExec would before a +/// sort preserving merge +fn partition_keys(keys: &[u64]) -> Vec> { + let mut partitions: Vec> = (0..NUM_PARTITIONS) + .map(|_| Vec::with_capacity(ROWS_PER_PARTITION)) + .collect(); + for (i, chunk) in keys.chunks(BATCH_SIZE).enumerate() { + partitions[i % NUM_PARTITIONS].extend_from_slice(chunk); + } + for partition in &mut partitions { + partition.sort_unstable(); + } + partitions +} + +fn key_columns(key_type: &str, keys: &[u64]) -> Vec<(String, ArrayRef)> { + let as_string = || { + Arc::new(StringArray::from_iter_values( + keys.iter().map(|k| format!("{k:012}")), + )) as ArrayRef + }; + match key_type { + "u64" => vec![( + "key0".to_string(), + Arc::new(UInt64Array::from(keys.to_vec())) as _, + )], + "string" => vec![("key0".to_string(), as_string())], + // Two sort columns force the row-based (normalized key) cursor + "complex" => vec![ + ( + "key0".to_string(), + Arc::new(UInt64Array::from_iter_values(keys.iter().map(|k| k / 8))) as _, + ), + ("key1".to_string(), as_string()), + ], + _ => unreachable!(), + } +} + +fn create_case( + ordering: &str, + key_type: &str, + payload_width: usize, +) -> (Vec>, SchemaRef, LexOrdering) { + let partitions = partition_keys(&generate_keys(ordering)) + .into_iter() + .map(|keys| { + keys.chunks(BATCH_SIZE) + .map(|chunk| { + let mut columns = key_columns(key_type, chunk); + // All payload columns share the same buffer, so wide + // payloads don't blow up memory + let payload = Arc::new(UInt64Array::from(chunk.to_vec())) as ArrayRef; + for i in 0..payload_width { + columns.push((format!("col{i}"), Arc::clone(&payload))); + } + RecordBatch::try_from_iter(columns).unwrap() + }) + .collect::>() + }) + .collect::>(); + + let schema = partitions[0][0].schema(); + let sort_order = LexOrdering::new( + schema + .fields() + .iter() + .filter(|field| field.name().starts_with("key")) + .map(|field| { + PhysicalSortExpr::new( + col(field.name(), &schema).unwrap(), + SortOptions::default(), + ) + }), + ) + .unwrap(); + + (partitions, schema, sort_order) +} + +fn bench_spm_data_patterns(c: &mut Criterion) { + let rt = tokio::runtime::Runtime::new().unwrap(); + let task_ctx = Arc::new(TaskContext::default()); + + for ordering in ORDERINGS { + for key_type in KEY_TYPES { + for payload_width in PAYLOAD_WIDTHS { + let (partitions, schema, sort_order) = + create_case(ordering, key_type, payload_width); + + c.bench_function( + &format!("spm/{ordering}/{key_type}/payload_{payload_width}"), + |b| { + b.iter(|| { + let exec = TestMemoryExec::try_new_exec( + &partitions, + Arc::clone(&schema), + None, + ) + .unwrap(); + let merge = Arc::new(SortPreservingMergeExec::new( + sort_order.clone(), + exec, + )); + rt.block_on(drain( + execute_stream(merge, Arc::clone(&task_ctx)).unwrap(), + )) + }) + }, + ); + } + } + } +} + +criterion_group!( + benches, + bench_merge_sorted_preserving, + bench_spm_data_patterns +); criterion_main!(benches); diff --git a/datafusion/physical-plan/src/sorts/builder.rs b/datafusion/physical-plan/src/sorts/builder.rs index 75eb2ff980325..274f182b83981 100644 --- a/datafusion/physical-plan/src/sorts/builder.rs +++ b/datafusion/physical-plan/src/sorts/builder.rs @@ -129,6 +129,23 @@ impl BatchBuilder { &self.schema } + /// Consume the rows of `stream_idx`'s current batch that have not been + /// pushed via [`Self::push_row`] yet, as a zero-copy slice. + /// + /// Also releases batches that are no longer referenced: skip-heavy merges + /// never reach [`Self::build_record_batch`], which otherwise does that. + pub fn take_remaining_rows_of(&mut self, stream_idx: usize) -> Option { + let cursor = &mut self.cursors[stream_idx]; + let (_, batch) = &self.batches[cursor.batch_idx]; + let remaining = batch.num_rows() - cursor.row_idx; + let slice = (remaining > 0).then(|| batch.slice(cursor.row_idx, remaining)); + cursor.row_idx = batch.num_rows(); + if self.indices.is_empty() { + self.release_unreferenced_batches(); + } + slice + } + /// Try to interleave all columns using the given index slice. fn try_interleave_columns( &self, @@ -164,28 +181,41 @@ impl BatchBuilder { // hot path. The retention is bounded and short-lived since leftover // rows are drained over subsequent polls. if self.indices.is_empty() { - // New cursors are only created once the previous cursor for the stream - // is finished. This means all remaining rows from all but the last batch - // for each stream have been yielded to the newly created record batch - // - // We can therefore drop all but the last batch for each stream - let mut batch_idx = 0; - let mut retained = 0; - self.batches.retain(|(stream_idx, batch)| { - let stream_cursor = &mut self.cursors[*stream_idx]; - let retain = stream_cursor.batch_idx == batch_idx; - batch_idx += 1; - - if retain { - stream_cursor.batch_idx = retained; - retained += 1; - } else { - self.batches_mem_used -= get_record_batch_memory_size(batch); - } - retain - }); + self.release_unreferenced_batches(); } + RecordBatch::try_new(Arc::clone(&self.schema), columns).map_err(Into::into) + } + + /// Drop buffered batches no stream cursor references anymore and release + /// their memory back to the pool. + /// + /// Must only be called when `indices` is empty, since indices reference + /// batches by position. + fn release_unreferenced_batches(&mut self) { + assert_eq!(self.indices.len(), 0); + + // New cursors are only created once the previous cursor for the stream + // is finished. This means all remaining rows from all but the last batch + // for each stream have been yielded to the newly created record batch + // + // We can therefore drop all but the last batch for each stream + let mut batch_idx = 0; + let mut retained = 0; + self.batches.retain(|(stream_idx, batch)| { + let stream_cursor = &mut self.cursors[*stream_idx]; + let retain = stream_cursor.batch_idx == batch_idx; + batch_idx += 1; + + if retain { + stream_cursor.batch_idx = retained; + retained += 1; + } else { + self.batches_mem_used -= get_record_batch_memory_size(batch); + } + retain + }); + // Release excess memory back to the pool, but never shrink below // initial_reservation to maintain the anti-starvation guarantee // for the merge phase. @@ -193,8 +223,6 @@ impl BatchBuilder { if self.reservation.size() > target { self.reservation.shrink(self.reservation.size() - target); } - - RecordBatch::try_new(Arc::clone(&self.schema), columns).map_err(Into::into) } /// Drains the in_progress row indexes, and builds a new RecordBatch from them diff --git a/datafusion/physical-plan/src/sorts/cursor.rs b/datafusion/physical-plan/src/sorts/cursor.rs index 8991922779d4a..58a4dff2bb384 100644 --- a/datafusion/physical-plan/src/sorts/cursor.rs +++ b/datafusion/physical-plan/src/sorts/cursor.rs @@ -43,8 +43,18 @@ pub trait CursorValues { fn eq_to_previous(cursor: &Self, idx: usize) -> bool; /// Returns comparison of `l[l_idx]` and `r[r_idx]` + /// + /// Only called with the cursors' *current* offsets, so caching + /// implementations may serve it from their cache. fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering; + /// Returns comparison of `l[l_idx]` and `r[r_idx]` for *arbitrary* + /// indices, unlike [`Self::compare`]. Cold path: caching implementations + /// must override this to index directly. + fn compare_at(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering { + Self::compare(l, l_idx, r, r_idx) + } + /// Notifies the values that the owning [`Cursor`] moved to `offset` (always /// `< len()`), so caching implementations can refresh the value(s) read by /// the hot comparisons. Default no-op (e.g. byte/row cursors don't benefit). @@ -116,6 +126,23 @@ impl Cursor { t } + /// Returns true if no rows of this cursor were consumed yet + pub fn is_at_start(&self) -> bool { + self.offset == 0 + } + + /// Compare this cursor's *last* row to `other`'s current row. Since the + /// values are sorted, `Less` means every remaining row of this cursor + /// sorts before all of `other`'s remaining rows. + pub fn compare_last_to(&self, other: &Self) -> Ordering { + T::compare_at( + &self.values, + self.values.len() - 1, + &other.values, + other.offset, + ) + } + pub fn is_eq_to_prev_one(&self, prev_cursor: Option<&Cursor>) -> bool { if self.offset > 0 { self.is_eq_to_prev_row() @@ -290,6 +317,11 @@ impl CursorValues for PrimitiveValues { l.current.compare(r.current) } + fn compare_at(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering { + // Arbitrary indices (cold path), so index directly instead of using the cache. + l.values[l_idx].compare(r.values[r_idx]) + } + #[inline(always)] fn set_offset(&mut self, offset: usize) { // The caller (`Cursor::advance`) guarantees `offset < len`; inlined, that @@ -518,6 +550,26 @@ impl CursorValues for ArrayValues { } } + fn compare_at(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering { + // Same null handling as `compare`, but delegates to the inner + // `compare_at` so arbitrary indices work with caching values. + match (l.is_null(l_idx), r.is_null(r_idx)) { + (true, true) => Ordering::Equal, + (true, false) => match l.options.nulls_first { + true => Ordering::Less, + false => Ordering::Greater, + }, + (false, true) => match l.options.nulls_first { + true => Ordering::Greater, + false => Ordering::Less, + }, + (false, false) => match l.options.descending { + true => T::compare_at(&r.values, r_idx, &l.values, l_idx), + false => T::compare_at(&l.values, l_idx, &r.values, r_idx), + }, + } + } + #[inline(always)] fn set_offset(&mut self, offset: usize) { // Forward to the wrapped values (e.g. caching `PrimitiveValues`). diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 4583d19e91061..1b9e3ac3cf156 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -18,6 +18,7 @@ //! Merge that deals with an arbitrary size of streaming inputs. //! This is an order-preserving merge. +use std::cmp::Ordering; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll, ready}; @@ -153,6 +154,10 @@ pub(crate) struct SortPreservingMergeStream { /// This vector contains the indices of the partitions that have not started emitting yet. uninitiated_partitions: Vec, + + /// A whole winner batch that was proven to sort before every other + /// stream's head, waiting to be emitted after the in-progress rows are flushed. + pending_skipped_batch: Option, } impl SortPreservingMergeStream { @@ -186,6 +191,7 @@ impl SortPreservingMergeStream { produced: 0, uninitiated_partitions: (0..stream_count).collect(), enable_round_robin_tie_breaker, + pending_skipped_batch: None, } } @@ -236,6 +242,12 @@ impl SortPreservingMergeStream { } return Poll::Ready(None); } + + if let Some(batch) = self.pending_skipped_batch.take() { + self.produced += batch.num_rows(); + return Poll::Ready(Some(Ok(batch))); + } + // Once all partitions have set their corresponding cursors for the loser tree, // we skip the following block. Until then, this function may be called multiple // times and can return Poll::Pending if any partition returns Poll::Pending. @@ -310,6 +322,44 @@ impl SortPreservingMergeStream { } let stream_idx = self.loser_tree[0]; + + // Full-batch skip: + // when the winner sits at the start of a fresh batch whose *last* row sorts before every other stream's + // current row, we skip per-row tree updates and just emit it. + // Checked at most once per input batch so we don't add a lot of overhead when + // the data is not sorted at all + // + // This might produce batches that do not equal to the configured batch_size, + // this is something we are willing to sacrifice + // + // TODO - support fetch + if self.fetch.is_none() + && self.cursors[stream_idx] + .as_ref() + .is_some_and(|cursor| cursor.is_at_start()) + && self.winner_batch_beats_all(stream_idx) + { + let batch = self + .in_progress + .take_remaining_rows_of(stream_idx) + .expect("winner cursor is at the start, so rows remain"); + // The batch is fully emitted: retire the cursor like `advance_cursors` does for a finished one + self.prev_cursors[stream_idx] = self.cursors[stream_idx].take(); + self.loser_tree_adjusted = false; + + if self.in_progress.is_empty() { + self.produced += batch.num_rows(); + return Poll::Ready(Some(Ok(batch))); + } + // Flush already-merged rows first + // they all sort before the skipped batch + // + // No need to reserve memory for holding this batch: it is a + // zero-copy slice of a batch the builder still buffers which also reserve memory for + self.pending_skipped_batch = Some(batch); + return Poll::Ready(self.emit_in_progress_batch().transpose()); + } + if self.advance_cursors(stream_idx) { self.loser_tree_adjusted = false; self.in_progress.push_row(stream_idx); @@ -327,6 +377,46 @@ impl SortPreservingMergeStream { } } + /// Returns `true` if every remaining row of `winner`'s current batch + /// sorts before every other stream's current row, using the same index + /// tie-breaking as [`Self::is_gt`]. Exhausted streams cannot compete. + fn winner_batch_beats_all(&self, winner: usize) -> bool { + let Some(winner_cursor) = &self.cursors[winner] else { + return false; + }; + + // Only the streams that lost *directly* to the winner can hold the + // overall runner-up, and those are exactly the losers stored on the + // winner's leaf-to-root path (every other stream lost to one of them + // transitively), so `O(log k)` comparisons suffice instead of `O(k)`. + let mut node = self.lt_leaf_node_index(winner); + while node != 0 { + let challenger = self.loser_tree[node]; + if let Some(other) = &self.cursors[challenger] { + let beats = match winner_cursor.compare_last_to(other) { + Ordering::Less => true, + Ordering::Equal => { + if self.enable_round_robin_tie_breaker { + // true if in round-robin because the winner already decided in a + // round-robin fashion + true + } else { + // Keep sort stable + winner < challenger + } + } + Ordering::Greater => false, + }; + if !beats { + return false; + } + } + node = self.lt_parent_node_index(node); + } + + true + } + /// For the given partition, updates the poll count. If the current value is the same /// of the previous value, it increases the count by 1; otherwise, it is reset as 0. fn update_poll_count_on_the_same_value(&mut self, partition_idx: usize) { @@ -534,15 +624,15 @@ impl SortPreservingMergeStream { if self.enable_round_robin_tie_breaker { match (&self.cursors[winner], &self.cursors[challenger]) { (Some(ac), Some(bc)) => match ac.cmp(bc) { - std::cmp::Ordering::Equal => { + Ordering::Equal => { self.handle_tie(cmp_node, &mut winner, challenger); } - std::cmp::Ordering::Greater => { + Ordering::Greater => { // Ends of tie breaker self.round_robin_tie_breaker_mode = false; self.update_winner(cmp_node, &mut winner, challenger); } - std::cmp::Ordering::Less => { + Ordering::Less => { // Ends of tie breaker self.round_robin_tie_breaker_mode = false; } diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 77a7d8f8f2e11..32f63ce6a2d72 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -442,7 +442,7 @@ mod tests { ArrayRef, Int32Array, Int64Array, RecordBatch, StringArray, TimestampNanosecondArray, }; - use arrow::compute::SortOptions; + use arrow::compute::{SortOptions, concat_batches}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion_common::test_util::batches_to_string; use datafusion_common::{assert_batches_eq, exec_err}; @@ -776,9 +776,9 @@ mod tests { context: Arc, ) -> RecordBatch { let merge = Arc::new(SortPreservingMergeExec::new(sort, input)); - let mut result = collect(merge, context).await.unwrap(); - assert_eq!(result.len(), 1); - result.remove(0) + let result = collect(merge, context).await.unwrap(); + // the merge may split the output into several batches, so concat first + concat_batches(&result[0].schema(), &result).unwrap() } async fn partition_sort( @@ -1139,10 +1139,10 @@ mod tests { .with_reservation(reservation) .build()?; - let mut merged = common::collect(merge_stream).await.unwrap(); + let merged = common::collect(merge_stream).await.unwrap(); - assert_eq!(merged.len(), 1); - let merged = merged.remove(0); + // the merge may split the output into several batches, so concat first + let merged = concat_batches(&merged[0].schema(), &merged).unwrap(); let basic = basic_sort(batches, sort.clone(), Arc::clone(&task_ctx)).await; let basic = arrow::util::pretty::pretty_format_batches(&[basic]) @@ -1557,8 +1557,16 @@ mod tests { .with_reservation(reservation) .build()?; - let first = merge_stream.next().await.unwrap(); - assert!(first.is_err(), "expected merge stream to surface the error"); + // the merge may emit already-proven-sorted batches before polling the + // failing stream again, so consume until the error surfaces + let mut saw_error = false; + while let Some(item) = merge_stream.next().await { + if item.is_err() { + saw_error = true; + break; + } + } + assert!(saw_error, "expected merge stream to surface the error"); assert!( merge_stream.next().await.is_none(), "merge stream yielded data after returning an error" diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index ade24ff0534ff..67f5ff7f0cb3d 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -265,3 +265,330 @@ impl<'a> StreamingMergeBuilder<'a> { ))) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::common::collect; + use crate::memory::MemoryStream; + use arrow::array::{Int32Array, RecordBatch}; + use arrow::datatypes::Int32Type; + use arrow_schema::{DataType, Field, Schema}; + use datafusion_execution::memory_pool::GreedyMemoryPool; + use datafusion_physical_expr::expressions::col; + use datafusion_physical_expr_common::metrics::ExecutionPlanMetricsSet; + use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; + use std::sync::Arc; + + /// Merge the given streams (each a list of batches of nullable `sort_key` + /// values, nulls sorted first) and return the merged values. + async fn merge_nullable_i32_streams( + inputs: Vec>>>, + batch_size: usize, + fetch: Option, + ) -> Vec>> { + merge_nullable_i32_streams_with_pool( + inputs, + batch_size, + fetch, + Arc::new(UnboundedMemoryPool::default()), + ) + .await + } + + async fn merge_nullable_i32_streams_with_pool( + inputs: Vec>>>, + batch_size: usize, + fetch: Option, + mem_pool: Arc, + ) -> Vec>> { + let schema = Arc::new(Schema::new(vec![Field::new( + "sort_key", + DataType::Int32, + true, + )])); + + let streams = inputs + .into_iter() + .map(|batches| { + let batches = batches + .into_iter() + .map(|values| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(values))], + ) + .unwrap() + }) + .collect::>(); + Box::pin( + MemoryStream::try_new(batches, Arc::clone(&schema), None).unwrap(), + ) as SendableRecordBatchStream + }) + .collect::>(); + + let merged = StreamingMergeBuilder::new() + .with_streams(streams) + .with_batch_size(batch_size) + .with_schema(Arc::clone(&schema)) + .with_reservation( + MemoryConsumer::new("merge stream mock memory").register(&mem_pool), + ) + .with_expressions( + &([ + PhysicalSortExpr::new_default(col("sort_key", &schema).unwrap()) + .asc(), + ] + .into()), + ) + .with_metrics(BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0)) + .with_fetch(fetch) + .build() + .unwrap(); + + let output = collect(merged).await.unwrap(); + + output + .iter() + .map(|b| b.column(0).as_primitive::().iter().collect()) + .collect() + } + + /// Same as [`merge_nullable_i32_streams`] for non-nullable values. + async fn merge_i32_streams( + inputs: Vec>>, + batch_size: usize, + fetch: Option, + ) -> Vec> { + let inputs = inputs + .into_iter() + .map(|stream| { + stream + .into_iter() + .map(|batch| batch.into_iter().map(Some).collect()) + .collect() + }) + .collect(); + + merge_nullable_i32_streams(inputs, batch_size, fetch) + .await + .into_iter() + .map(|batch| batch.into_iter().map(|v| v.unwrap()).collect()) + .collect() + } + + fn flatten_vec(input: Vec>) -> Vec { + input.into_iter().flatten().collect() + } + + #[tokio::test] + async fn test_merge_non_overlapping_batches() { + // streams take turns holding the smallest batch + let streams = vec![ + vec![ + vec![1, 2, 3, 4, 5], + vec![16, 17, 18, 19, 20], + vec![31, 32, 33], + ], + vec![ + vec![6, 7, 8, 9, 10], + vec![21, 22, 23, 24, 25], + vec![34, 35, 36], + ], + vec![ + vec![11, 12, 13, 14, 15], + vec![26, 27, 28, 29, 30], + vec![37, 38], + ], + ]; + + let output = merge_i32_streams(streams, 5, None).await; + + assert_eq!(flatten_vec(output), (1..=38).collect::>()); + } + + #[tokio::test] + async fn test_merge_full_batch_release_memory() { + // The pool only fits a few batches of the ~400KB total input: + // if skipped batches were not released as the merge progresses, the reservation would exhaust the pool. + const ROWS: i32 = 1024; // ~4KB per batch + const BATCHES_PER_STREAM: i32 = 50; + + let batch_of = |start: i32| (start..start + ROWS).map(Some).collect::>(); + let streams = (0..2) + .map(|stream| { + (0..BATCHES_PER_STREAM) + .map(|batch| batch_of((batch * 2 + stream) * ROWS)) + .collect() + }) + .collect(); + + let output = merge_nullable_i32_streams_with_pool( + streams, + ROWS as usize, + None, + Arc::new(GreedyMemoryPool::new(64 * 1024)), + ) + .await; + + let expected = (0..2 * BATCHES_PER_STREAM * ROWS) + .map(Some) + .collect::>(); + assert_eq!(flatten_vec(output), expected); + } + + #[tokio::test] + async fn test_merge_interleaved_rows() { + // every batch overlaps batches of the other streams + let streams = vec![ + vec![vec![1, 4, 7], vec![10, 13, 16]], + vec![vec![2, 5, 8], vec![11, 14, 17]], + vec![vec![3, 6, 9], vec![12, 15, 18]], + ]; + + let output = merge_i32_streams(streams, 4, None).await; + + assert_eq!(flatten_vec(output), (1..=18).collect::>()); + } + + #[tokio::test] + async fn test_merge_some_rows_from_other_streams_then_a_whole_batch_from_one() { + // one row from stream 0, two from stream 1 (batch_size 4 is still + // not filled), then stream 2's entire batch [4, 5, 6, 7] sorts before + // everything the other streams have left (8 and 11) + let streams = vec![ + vec![vec![1, 8, 9, 10], vec![17, 18, 19, 20]], + vec![vec![2, 3, 11, 12], vec![21, 22, 23, 24]], + vec![vec![4, 5, 6, 7], vec![13, 14, 15, 16]], + ]; + + let output = merge_i32_streams(streams, 4, None).await; + + assert_eq!(flatten_vec(output), (1..=24).collect::>()); + } + + #[tokio::test] + async fn test_merge_equal_values_across_batch_and_stream_boundaries() { + // the value 3 ends stream 0's first batch, starts its second + // one, and appears in every other stream as well + let streams = vec![ + vec![vec![1, 2, 3], vec![3, 3, 4]], + vec![vec![3, 4, 5], vec![5, 6, 7]], + vec![vec![3, 3, 3], vec![7, 7, 7]], + ]; + + let output = merge_i32_streams(streams, 3, None).await; + + assert_eq!( + flatten_vec(output), + vec![1, 2, 3, 3, 3, 3, 3, 3, 3, 4, 4, 5, 5, 6, 7, 7, 7, 7] + ); + } + + #[tokio::test] + async fn test_merge_with_nulls() { + // nulls sort first; stream 0's first batch is all nulls, so its last + // row (null) ties with the null heads of the other streams, and nulls + // also compare against non-null values on every batch boundary + let streams = vec![ + vec![vec![None, None], vec![Some(1), Some(2)]], + vec![vec![None, Some(3)], vec![Some(4), Some(5)]], + vec![vec![None, Some(6)], vec![Some(7), Some(8)]], + ]; + + let output = merge_nullable_i32_streams(streams, 4, None).await; + + assert_eq!( + flatten_vec(output), + vec![ + None, + None, + None, + None, + Some(1), + Some(2), + Some(3), + Some(4), + Some(5), + Some(6), + Some(7), + Some(8) + ] + ); + } + + #[tokio::test] + async fn test_merge_streams_of_different_lengths() { + // stream 1 ends first, then stream 2, stream 0 finishes alone + let streams = vec![ + vec![ + vec![1, 2], + vec![3, 4], + vec![5, 6], + vec![7, 8], + vec![20, 21, 22], + ], + vec![vec![0, 9]], + vec![vec![4, 5], vec![10, 11, 12]], + ]; + + let output = merge_i32_streams(streams, 4, None).await; + + assert_eq!( + flatten_vec(output), + vec![0, 1, 2, 3, 4, 4, 5, 5, 6, 7, 8, 9, 10, 11, 12, 20, 21, 22] + ); + } + + #[tokio::test] + async fn test_merge_batch_size_smaller_than_input_batches() { + let streams = vec![ + vec![vec![1, 4, 7, 10], vec![13, 14, 15, 16]], + vec![vec![2, 5, 8, 11]], + vec![vec![3, 6, 9, 12]], + ]; + + let output = merge_i32_streams(streams, 2, None).await; + + assert_eq!(flatten_vec(output), (1..=16).collect::>()); + } + + #[tokio::test] + async fn test_merge_batch_size_larger_than_entire_input() { + let streams = vec![ + vec![vec![1, 4, 7, 10], vec![13, 14, 15, 16]], + vec![vec![2, 5, 8, 11]], + vec![vec![3, 6, 9, 12]], + ]; + + let output = merge_i32_streams(streams, 8192, None).await; + + assert_eq!(flatten_vec(output), (1..=16).collect::>()); + } + + #[tokio::test] + async fn test_merge_fetch_smaller_than_input() { + let streams = vec![ + vec![vec![1, 4, 7], vec![10, 11, 12]], + vec![vec![2, 5, 8]], + vec![vec![3, 6, 9]], + ]; + + let output = merge_i32_streams(streams, 4, Some(5)).await; + + assert_eq!(flatten_vec(output), vec![1, 2, 3, 4, 5]); + } + + #[tokio::test] + async fn test_merge_fetch_larger_than_input() { + let streams = vec![ + vec![vec![1, 4, 7], vec![10, 11, 12]], + vec![vec![2, 5, 8]], + vec![vec![3, 6, 9]], + ]; + + let output = merge_i32_streams(streams, 4, Some(100)).await; + + assert_eq!(flatten_vec(output), (1..=12).collect::>()); + } +}