diff --git a/datafusion/physical-plan/benches/sort_preserving_merge.rs b/datafusion/physical-plan/benches/sort_preserving_merge.rs index 76ebf230a30e0..a2858e70c4af6 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,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 = (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. + // + // 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. + for partitions in [2, 4] { + for (label, round_robin) in [("on", true), ("off", false)] { + c.bench_function( + &format!( + "bench_merge_tied_keys_slow_producers/{partitions}_partitions/tie_breaker_{label}" + ), + |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), + ) + .with_round_robin_tie_breaker(round_robin) + .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 bdfdbcf38b839..d08bd55de91ca 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/merge_trace.rs b/datafusion/physical-plan/src/sorts/merge_trace.rs new file mode 100644 index 0000000000000..95f9084e58a64 --- /dev/null +++ b/datafusion/physical-plan/src/sorts/merge_trace.rs @@ -0,0 +1,441 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Snapshot tests for *which input the merge drains, and in what order*. +//! +//! The round-robin tie breaker in [`SortPreservingMergeStream`] never changes +//! the merged output: rows with equal sort keys are interchangeable, so the +//! sorted values are identical whether it is on or off. All it changes is +//! which tied input gets consumed next, and therefore which upstream producers +//! keep making progress. Ordinary correctness tests are blind to that, which +//! makes the behavior easy to regress silently. +//! +//! This module drives a merge over hand-written inputs and records an +//! interleaved trace of everything an outside observer can see: +//! +//! * `poll S` — the merge asked input `n` for another batch +//! * `row S` — the merge emitted a row that came from input `n` +//! +//! The traces are asserted with inline `insta` snapshots, so a change in poll +//! order shows up as a snapshot diff instead of passing unnoticed. To accept +//! an intended change, run the tests with `INSTA_FORCE_UPDATE=1` and apply the +//! result with `cargo insta accept`, rather than hand-editing the expected +//! text -- the point is that changing it is a deliberate act. +//! +//! [`SortPreservingMergeStream`]: crate::sorts::merge::SortPreservingMergeStream + +use std::fmt::Write as _; +use std::pin::Pin; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; + +use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet}; +use crate::sorts::streaming_merge::StreamingMergeBuilder; +use crate::stream::RecordBatchStreamAdapter; +use crate::{RecordBatchStream, SendableRecordBatchStream}; + +use arrow::array::{AsArray, Int32Array, RecordBatch}; +use arrow::datatypes::{DataType, Field, Int32Type, Schema, SchemaRef}; +use datafusion_common::Result; +use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr, expressions::col}; + +use futures::{Stream, StreamExt}; +use insta::assert_snapshot; + +/// One event observed while driving a merge. +#[derive(Debug)] +enum TraceEvent { + /// The merge polled input `partition`, which returned a batch of `keys`. + Batch { partition: usize, keys: Vec }, + /// The merge polled input `partition`, which reported end of stream. + Eof { partition: usize }, + /// The merge emitted an output row that originated in `partition`. + Row { partition: usize, key: i32 }, +} + +/// Shared, append-only log of [`TraceEvent`]s. +/// +/// The input streams append to it from inside the merge's `poll_next`, and the +/// driver appends emitted rows in between, which is what makes the two kinds +/// of event interleave in issue order. +type TraceLog = Arc>>; + +/// Wraps one input of the merge and logs every poll that resolves. +/// +/// `Poll::Pending` is not logged: these inputs are always immediately ready, +/// so a pending poll never happens. +struct TracingStream { + partition: usize, + inner: SendableRecordBatchStream, + log: TraceLog, +} + +impl Stream for TracingStream { + type Item = Result; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + let partition = self.partition; + let polled = self.inner.poll_next_unpin(cx); + match &polled { + Poll::Ready(Some(Ok(batch))) => { + let keys = batch + .column(0) + .as_primitive::() + .values() + .to_vec(); + self.log + .lock() + .unwrap() + .push(TraceEvent::Batch { partition, keys }); + } + Poll::Ready(None) => { + self.log.lock().unwrap().push(TraceEvent::Eof { partition }); + } + Poll::Ready(Some(Err(_))) | Poll::Pending => {} + } + polled + } +} + +impl RecordBatchStream for TracingStream { + fn schema(&self) -> SchemaRef { + self.inner.schema() + } +} + +/// `(key, source partition)` pairs; `key` is the sort key, the partition index +/// rides along so emitted rows can be attributed back to their input. +fn trace_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("src", DataType::Int32, false), + ])) +} + +/// Builds one input stream: `batches` is a list of batches, each a list of +/// sort keys. All rows are tagged with `partition`. +fn input_stream(partition: usize, batches: &[&[i32]], log: &TraceLog) -> TracingStream { + let schema = trace_schema(); + let batches: Vec> = batches + .iter() + .map(|keys| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(keys.to_vec())), + Arc::new(Int32Array::from(vec![partition as i32; keys.len()])), + ], + ) + .map_err(Into::into) + }) + .collect(); + + TracingStream { + partition, + inner: Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::iter(batches), + )), + log: Arc::clone(log), + } +} + +/// Merges `inputs` (partition -> batch -> sort keys) and returns the rendered +/// trace of polls and emitted rows. +/// +/// The merge runs with `batch_size = 1` so every output row is handed back +/// individually and can be slotted into the log at the point it was produced. +/// Output batching is independent of the loser tree and the tie breaker, so +/// this makes the trace finer-grained without changing what is being traced. +async fn trace_merge(inputs: &[&[&[i32]]], round_robin_tie_breaker: bool) -> String { + let schema = trace_schema(); + let log: TraceLog = Arc::new(Mutex::new(Vec::new())); + + let streams: Vec = inputs + .iter() + .enumerate() + .map(|(partition, batches)| { + Box::pin(input_stream(partition, batches, &log)) as SendableRecordBatchStream + }) + .collect(); + + let ordering: LexOrdering = + [PhysicalSortExpr::new_default(col("key", &schema).unwrap())].into(); + + let mut merged = StreamingMergeBuilder::new() + .with_streams(streams) + .with_schema(Arc::clone(&schema)) + .with_expressions(&ordering) + .with_metrics(BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0)) + .with_batch_size(1) + .with_bypass_mempool() + .with_round_robin_tie_breaker(round_robin_tie_breaker) + .build() + .unwrap(); + + while let Some(batch) = merged.next().await { + let batch = batch.unwrap(); + let keys = batch.column(0).as_primitive::(); + let sources = batch.column(1).as_primitive::(); + let mut log = log.lock().unwrap(); + for row in 0..batch.num_rows() { + log.push(TraceEvent::Row { + partition: sources.value(row) as usize, + key: keys.value(row), + }); + } + } + + render(inputs, round_robin_tie_breaker, &log.lock().unwrap()) +} + +/// Renders the inputs and the recorded events as the snapshot text. +fn render( + inputs: &[&[&[i32]]], + round_robin_tie_breaker: bool, + events: &[TraceEvent], +) -> String { + let mut out = String::new(); + writeln!(out, "round_robin_tie_breaker = {round_robin_tie_breaker}").unwrap(); + writeln!(out, "inputs:").unwrap(); + for (partition, batches) in inputs.iter().enumerate() { + let batches = batches + .iter() + .map(|keys| format!("[{}]", join(keys))) + .collect::>() + .join(" "); + writeln!(out, " S{partition}: {batches}").unwrap(); + } + writeln!(out, "trace:").unwrap(); + for event in events { + match event { + TraceEvent::Batch { partition, keys } => { + writeln!(out, " poll S{partition} -> [{}]", join(keys)).unwrap() + } + TraceEvent::Eof { partition } => { + writeln!(out, " poll S{partition} -> done").unwrap() + } + TraceEvent::Row { partition, key } => { + writeln!(out, " row S{partition} key={key}").unwrap() + } + } + } + out +} + +fn join(keys: &[i32]) -> String { + keys.iter() + .map(|k| k.to_string()) + .collect::>() + .join(",") +} + +/// Traces the same inputs with the tie breaker enabled and disabled, so a +/// single snapshot shows what the tie breaker changes. +async fn trace_both(inputs: &[&[&[i32]]]) -> String { + format!( + "{}\n{}", + trace_merge(inputs, true).await, + trace_merge(inputs, false).await + ) +} + +/// With every key tied, the tie breaker must hand rows out in strict +/// round-robin order; with it off, the lowest-index input is drained first and +/// the other producer is left idle. +#[tokio::test] +async fn test_tied_keys_alternate_between_inputs() { + let inputs: &[&[&[i32]]] = &[&[&[1, 1, 1]], &[&[1, 1, 1]]]; + + assert_snapshot!(trace_both(inputs).await, @r" + round_robin_tie_breaker = true + inputs: + S0: [1,1,1] + S1: [1,1,1] + trace: + poll S0 -> [1,1,1] + poll S1 -> [1,1,1] + row S0 key=1 + row S1 key=1 + row S0 key=1 + row S1 key=1 + row S0 key=1 + poll S0 -> done + row S1 key=1 + poll S1 -> done + + round_robin_tie_breaker = false + inputs: + S0: [1,1,1] + S1: [1,1,1] + trace: + poll S0 -> [1,1,1] + poll S1 -> [1,1,1] + row S0 key=1 + row S0 key=1 + row S0 key=1 + poll S0 -> done + row S1 key=1 + row S1 key=1 + row S1 key=1 + poll S1 -> done + "); +} + +/// Regression test for stale poll counts leaking across runs of ties. +/// +/// `S0` runs out of `1`s first, so the first run of ties ends with `S1` still +/// holding several rows that it drains on its own. The second run (key `2`) +/// is a fresh tie-breaker round and must alternate from its very first row. +/// +/// Before the fix, `is_poll_count_gt` compared the challenger's *raw* counter, +/// which still held the count `S1` accumulated during the key `1` run, and the +/// key `2` run started `S1 S0 S0 S1` instead of alternating. +#[tokio::test] +async fn test_poll_counts_reset_between_runs_of_ties() { + let inputs: &[&[&[i32]]] = &[&[&[1, 1], &[2, 2]], &[&[1, 1, 1, 1], &[2, 2]]]; + + assert_snapshot!(trace_both(inputs).await, @r" + round_robin_tie_breaker = true + inputs: + S0: [1,1] [2,2] + S1: [1,1,1,1] [2,2] + trace: + poll S0 -> [1,1] + poll S1 -> [1,1,1,1] + row S0 key=1 + row S1 key=1 + row S0 key=1 + poll S0 -> [2,2] + row S1 key=1 + row S1 key=1 + row S1 key=1 + poll S1 -> [2,2] + row S0 key=2 + row S1 key=2 + row S0 key=2 + poll S0 -> done + row S1 key=2 + poll S1 -> done + + round_robin_tie_breaker = false + inputs: + S0: [1,1] [2,2] + S1: [1,1,1,1] [2,2] + trace: + poll S0 -> [1,1] + poll S1 -> [1,1,1,1] + row S0 key=1 + row S0 key=1 + poll S0 -> [2,2] + row S1 key=1 + row S1 key=1 + row S1 key=1 + row S1 key=1 + poll S1 -> [2,2] + row S0 key=2 + row S0 key=2 + poll S0 -> done + row S1 key=2 + row S1 key=2 + poll S1 -> done + "); +} + +/// Without ties there is nothing to break: both settings must produce the same +/// poll order. A diff here means the tie breaker started interfering with the +/// ordinary loser-tree path. +#[tokio::test] +async fn test_distinct_keys_are_unaffected_by_the_tie_breaker() { + let inputs: &[&[&[i32]]] = &[&[&[1, 3, 5]], &[&[2, 4, 6]]]; + + assert_snapshot!(trace_both(inputs).await, @r" + round_robin_tie_breaker = true + inputs: + S0: [1,3,5] + S1: [2,4,6] + trace: + poll S0 -> [1,3,5] + poll S1 -> [2,4,6] + row S0 key=1 + row S1 key=2 + row S0 key=3 + row S1 key=4 + row S0 key=5 + poll S0 -> done + row S1 key=6 + poll S1 -> done + + round_robin_tie_breaker = false + inputs: + S0: [1,3,5] + S1: [2,4,6] + trace: + poll S0 -> [1,3,5] + poll S1 -> [2,4,6] + row S0 key=1 + row S1 key=2 + row S0 key=3 + row S1 key=4 + row S0 key=5 + poll S0 -> done + row S1 key=6 + poll S1 -> done + "); +} + +/// The tie breaker only runs at the root of the loser tree (`cmp_node == 1`), +/// so with four tied inputs it balances the two sub-tree winners rather than +/// all four producers: `S1` and `S3` still drain in one go. +/// +/// This is a known limitation, recorded here so that widening the tie breaker +/// to inner nodes shows up as a deliberate snapshot change. Note also that the +/// initial polls are not in partition order — `initialize_all_partitions` uses +/// `swap_remove`, which reorders the pending list. +#[tokio::test] +async fn test_tie_breaker_only_balances_the_root_comparison() { + let inputs: &[&[&[i32]]] = &[&[&[1, 1]], &[&[1, 1]], &[&[1, 1]], &[&[1, 1]]]; + + assert_snapshot!(trace_merge(inputs, true).await, @r" + round_robin_tie_breaker = true + inputs: + S0: [1,1] + S1: [1,1] + S2: [1,1] + S3: [1,1] + trace: + poll S0 -> [1,1] + poll S3 -> [1,1] + poll S2 -> [1,1] + poll S1 -> [1,1] + row S0 key=1 + row S2 key=1 + row S0 key=1 + poll S0 -> done + row S1 key=1 + row S1 key=1 + poll S1 -> done + row S2 key=1 + poll S2 -> done + row S3 key=1 + row S3 key=1 + poll S3 -> done + "); +} diff --git a/datafusion/physical-plan/src/sorts/mod.rs b/datafusion/physical-plan/src/sorts/mod.rs index ca8d4a4400c49..4af2a3629fa9c 100644 --- a/datafusion/physical-plan/src/sorts/mod.rs +++ b/datafusion/physical-plan/src/sorts/mod.rs @@ -20,6 +20,8 @@ mod builder; mod cursor; mod merge; +#[cfg(test)] +mod merge_trace; mod multi_level_merge; pub mod partial_sort; pub mod partitioned_topk;