diff --git a/datafusion/physical-plan/benches/sort_preserving_merge.rs b/datafusion/physical-plan/benches/sort_preserving_merge.rs index 76ebf230a30e..e68fe61a457e 100644 --- a/datafusion/physical-plan/benches/sort_preserving_merge.rs +++ b/datafusion/physical-plan/benches/sort_preserving_merge.rs @@ -16,7 +16,7 @@ // under the License. use arrow::{ - array::{ArrayRef, StringArray, UInt64Array}, + array::{ArrayRef, AsArray, StringArray, UInt64Array}, record_batch::RecordBatch, }; use arrow_schema::{SchemaRef, SortOptions}; @@ -193,5 +193,120 @@ fn bench_merge_sorted_preserving(c: &mut Criterion) { } } -criterion_group!(benches, bench_merge_sorted_preserving); +/// Merge inputs whose keys are mostly *tied* and whose producers do real work +/// per batch. +/// +/// `SortPreservingMergeExec` runs each input in its own task, buffered one +/// batch ahead (`spawn_buffered(_, 1)`). If the merge keeps draining a single +/// partition during a run of equal keys, that partition's producer becomes the +/// bottleneck while the others idle on their one buffered batch. The +/// round-robin tie breaker is meant to spread consumption across the tied +/// partitions so all producers stay busy. +fn bench_merge_tied_keys_slow_producers(c: &mut Criterion) { + use datafusion_execution::memory_pool::{ + MemoryConsumer, MemoryPool, UnboundedMemoryPool, + }; + use datafusion_physical_plan::common::spawn_buffered; + use datafusion_physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet}; + use datafusion_physical_plan::sorts::streaming_merge::StreamingMergeBuilder; + use datafusion_physical_plan::stream::RecordBatchStreamAdapter; + use futures::StreamExt; + + const ROWS: usize = 400_000; + const BATCH: usize = 8192; + const ROWS_PER_KEY: usize = 100_000; + + let schema: SchemaRef = Arc::new(arrow_schema::Schema::new(vec![ + arrow_schema::Field::new("key", arrow_schema::DataType::UInt64, false), + arrow_schema::Field::new("val", arrow_schema::DataType::UInt64, false), + ])); + let sort_order = LexOrdering::new(vec![PhysicalSortExpr::new( + col("key", &schema).unwrap(), + SortOptions::default(), + )]) + .unwrap(); + + // Every partition holds the same long runs of equal keys. + let batches: Vec = (0..ROWS.div_ceil(BATCH)) + .map(|b| { + let start = b * BATCH; + let n = BATCH.min(ROWS - start); + let keys = UInt64Array::from_iter_values( + (start..start + n).map(|i| (i / ROWS_PER_KEY) as u64), + ); + let vals = + UInt64Array::from_iter_values((start..start + n).map(|i| i as u64)); + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(keys), Arc::new(vals)], + ) + .unwrap() + }) + .collect(); + + /// Stand-in for an upstream operator: ~fixed CPU cost per batch. + fn produce(batch: RecordBatch) -> RecordBatch { + let vals = batch + .column(1) + .as_primitive::(); + let mut acc = 0u64; + for _ in 0..200 { + for v in vals.values() { + acc = acc.wrapping_mul(6364136223846793005).wrapping_add(*v); + } + } + std::hint::black_box(acc); + batch + } + + let rt = tokio::runtime::Runtime::new().unwrap(); + // With 2 inputs the root comparison is the whole tree, so the tie breaker + // balances all producers; with 4 it only balances the two sub-tree winners. + for partitions in [2, 4] { + c.bench_function( + &format!("bench_merge_tied_keys_slow_producers/{partitions}_partitions"), + |b| { + b.iter(|| { + rt.block_on(async { + let streams = (0..partitions) + .map(|_| { + let s = RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(batches.clone()) + .map(|b| Ok(produce(b))), + ); + spawn_buffered(Box::pin(s), 1) + }) + .collect(); + let pool: Arc = + Arc::new(UnboundedMemoryPool::default()); + let merged = StreamingMergeBuilder::new() + .with_streams(streams) + .with_schema(Arc::clone(&schema)) + .with_expressions(&sort_order) + .with_metrics(BaselineMetrics::new( + &ExecutionPlanMetricsSet::new(), + 0, + )) + .with_batch_size(BATCH) + .with_reservation( + MemoryConsumer::new("bench").register(&pool), + ) + .build() + .unwrap(); + datafusion_physical_plan::common::collect(merged) + .await + .unwrap(); + }) + }) + }, + ); + } +} + +criterion_group!( + benches, + bench_merge_sorted_preserving, + bench_merge_tied_keys_slow_producers +); criterion_main!(benches); diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index bdfdbcf38b83..d08bd55de91c 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -365,6 +365,21 @@ impl SortPreservingMergeStream { } } + /// Returns the poll count of `partition_idx` for the current tie-breaker + /// round. + /// + /// Poll counts are reset lazily by bumping `current_reset_epoch` (see + /// [`Self::reset_poll_counts`]), so a count written in an older epoch is + /// stale and reads as 0. + #[inline] + fn poll_count(&self, partition_idx: usize) -> usize { + if self.poll_reset_epochs[partition_idx] == self.current_reset_epoch { + self.num_of_polled_with_same_value[partition_idx] + } else { + 0 + } + } + /// 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) { @@ -457,11 +472,18 @@ impl SortPreservingMergeStream { } } + /// Returns `true` if partition `a` has been polled more often than `b` in + /// the current tie-breaker round, breaking equal counts by partition index. + /// + /// Both counts go through [`Self::poll_count`]: only the winner's count is + /// refreshed by [`Self::update_poll_count_on_the_same_value`] before this + /// is called, so the challenger's raw count may belong to an earlier round. #[inline] fn is_poll_count_gt(&self, a: usize, b: usize) -> bool { - let poll_a = self.num_of_polled_with_same_value[a]; - let poll_b = self.num_of_polled_with_same_value[b]; - poll_a.cmp(&poll_b).then_with(|| a.cmp(&b)).is_gt() + self.poll_count(a) + .cmp(&self.poll_count(b)) + .then_with(|| a.cmp(&b)) + .is_gt() } #[inline] diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index 81adad8e9ec8..f5ea8aadd616 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -281,7 +281,8 @@ mod tests { use super::*; - use arrow::array::{ArrayRef, RecordBatch}; + use arrow::array::{ArrayRef, AsArray, RecordBatch}; + use arrow::datatypes::{Field, Int32Type, Schema}; use arrow_schema::SortOptions; use datafusion_common::Result; use datafusion_execution::TaskContext; @@ -379,4 +380,78 @@ mod tests { Ok(()) } + + /// Merge streams of `(key, tag)` rows sorted on `key` with the round-robin + /// tie breaker enabled, returning the `tag` column of the output in order. + async fn merge_tags(streams: Vec>) -> Vec { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("tag", DataType::Int32, false), + ])); + let streams = streams + .into_iter() + .map(|rows| { + let (keys, tags): (Vec, Vec) = rows.into_iter().unzip(); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(keys)), + Arc::new(Int32Array::from(tags)), + ], + ) + .unwrap(); + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(vec![Ok(batch)]), + )) as SendableRecordBatchStream + }) + .collect(); + let sort: LexOrdering = + [PhysicalSortExpr::new_default(col("key", &schema).unwrap())].into(); + + let merged = StreamingMergeBuilder::new() + .with_streams(streams) + .with_schema(schema) + .with_expressions(&sort) + .with_metrics(BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0)) + .with_batch_size(1024) + .with_bypass_mempool() + .with_round_robin_tie_breaker(true) + .build() + .unwrap(); + + collect(merged) + .await + .unwrap() + .iter() + .flat_map(|b| b.column(1).as_primitive::().values().to_vec()) + .collect() + } + + /// The round-robin tie breaker must start every run of equal keys with a + /// clean slate: poll counts left over from an earlier run of ties must not + /// influence which stream wins the next one. + #[tokio::test] + async fn test_round_robin_tie_breaker_resets_poll_counts_between_tie_runs() { + // Stream 0 runs out of `1`s first, so the first tie run ends with + // stream 1 holding several unanswered rows. The second run (key `2`) + // must then alternate from its first row rather than let stream 1 + // "catch up" on the stale count stream 0 accumulated during run one. + let stream0: Vec<_> = std::iter::repeat_n((1, 0), 6) + .chain(std::iter::repeat_n((2, 0), 8)) + .collect(); + let stream1: Vec<_> = std::iter::repeat_n((1, 1), 12) + .chain(std::iter::repeat_n((2, 1), 8)) + .collect(); + + let tags = merge_tags(vec![stream0, stream1]).await; + + let expected: Vec = [0, 1] + .repeat(6) + .into_iter() + .chain(std::iter::repeat_n(1, 6)) + .chain([0, 1].repeat(8)) + .collect(); + assert_eq!(tags, expected); + } }