jayzhan211 commented on code in PR #24585:
URL: https://github.com/apache/datafusion/pull/24585#discussion_r3863385793
##########
datafusion/physical-plan/benches/sort_preserving_merge.rs:
##########
@@ -193,5 +193,130 @@ 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<RecordBatch> = (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::<arrow::datatypes::UInt64Type>();
+ 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.
+ //
+ // Each case is run with the tie breaker both enabled and disabled so the
+ // pair measures what the tie breaker actually buys on this workload: if
+ // the two ever converge, the balancing has regressed into the
+ // lowest-index-wins baseline.
Review Comment:
`benches/sort_preserving_merge.rs` now runs each partition count with
`.with_round_robin_tie_breaker(true)` and `(false)`, so the pair directly
measures what the tie breaker buys on this workload.
```
cargo bench -p datafusion-physical-plan --bench sort_preserving_merge -- \
bench_merge_tied_keys_slow_producers
```
| case | tie breaker on | tie breaker off | speedup |
| ------------ | -------------- | --------------- | ------- |
| 2 partitions | 82.8 ms | 150.8 ms | 1.82x |
| 4 partitions | 228.7 ms | 303.3 ms | 1.33x |
The gap is the thing to watch: if the enabled path ever regresses toward the
lowest-index-wins baseline, these two columns converge.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
