diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs index bc69ae5140831..673093edfdd00 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs @@ -211,7 +211,9 @@ impl ClassicPWMJStream { &mut self, cx: &mut std::task::Context<'_>, ) -> Poll>>> { - match ready!(self.streamed.poll_next_unpin(cx)) { + let next_batch = ready!(self.streamed.poll_next_unpin(cx)); + let _join_timer = self.join_metrics.join_time.timer(); + match next_batch { None => { // Release the streamed input pipeline's resources. let streamed_schema = self.streamed.schema(); @@ -266,6 +268,7 @@ impl ClassicPWMJStream { fn process_stream_batch( &mut self, ) -> Result>> { + let _join_timer = self.join_metrics.join_time.timer(); let buffered_side = self.buffered_side.try_as_ready_mut()?; let stream_batch = self.state.try_as_process_stream_batch_mut()?; @@ -312,6 +315,7 @@ impl ClassicPWMJStream { fn process_unmatched_buffered_batch( &mut self, ) -> Result>> { + let _join_timer = self.join_metrics.join_time.timer(); // Return early for `JoinType::Right` and `JoinType::Inner` if matches!(self.join_type, JoinType::Right | JoinType::Inner) { self.state = PiecewiseMergeJoinStreamState::Completed; @@ -416,7 +420,8 @@ impl Stream for ClassicPWMJStream { mut self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>, ) -> Poll> { - self.poll_next_impl(cx) + let poll = self.poll_next_impl(cx); + self.join_metrics.baseline.record_poll(poll) } } @@ -635,19 +640,24 @@ fn create_unmatched_batch( mod tests { use super::*; use crate::{ - ExecutionPlan, common, + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, common, joins::PiecewiseMergeJoinExec, - test::{TestMemoryExec, build_table_i32}, + metrics::MetricsSet, + stream::RecordBatchStreamAdapter, + test::{TestMemoryExec, assert_join_metrics, build_table_i32}, }; use arrow::array::{Date32Array, Date64Array}; use arrow_schema::{DataType, Field}; + use datafusion_common::instant::Instant; use datafusion_common::test_util::batches_to_string; + use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_execution::TaskContext; use datafusion_execution::config::SessionConfig; use datafusion_physical_expr::{PhysicalExpr, expressions::Column}; use futures::TryStreamExt; use insta::assert_snapshot; use std::sync::Arc; + use std::time::Duration; fn columns(schema: &Schema) -> Vec { schema.fields().iter().map(|f| f.name().clone()).collect() @@ -729,7 +739,7 @@ mod tests { on: (PhysicalExprRef, PhysicalExprRef), operator: Operator, join_type: JoinType, - ) -> Result<(Vec, Vec)> { + ) -> Result<(Vec, Vec, MetricsSet)> { join_collect_with_options(left, right, on, operator, join_type).await } @@ -739,14 +749,15 @@ mod tests { on: (PhysicalExprRef, PhysicalExprRef), operator: Operator, join_type: JoinType, - ) -> Result<(Vec, Vec)> { + ) -> Result<(Vec, Vec, MetricsSet)> { let task_ctx = Arc::new(TaskContext::default()); let join = join(left, right, on, operator, join_type)?; let columns = columns(&join.schema()); let stream = join.execute(0, task_ctx)?; let batches = common::collect(stream).await?; - Ok((columns, batches)) + let metrics = join.metrics().expect("metrics should be available"); + Ok((columns, batches, metrics)) } #[tokio::test] @@ -782,7 +793,7 @@ mod tests { Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, ); - let (_, batches) = + let (_, batches, metrics) = join_collect(left, right, on, Operator::Lt, JoinType::Inner).await?; assert_snapshot!(batches_to_string(&batches), @r" @@ -797,6 +808,8 @@ mod tests { | 3 | 1 | 9 | 10 | 2 | 70 | +----+----+----+----+----+----+ "); + + assert_join_metrics!(metrics, 6); Ok(()) } @@ -833,7 +846,7 @@ mod tests { Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, ); - let (_, batches) = + let (_, batches, metrics) = join_collect(left, right, on, Operator::Lt, JoinType::Inner).await?; assert_snapshot!(batches_to_string(&batches), @r" @@ -848,6 +861,8 @@ mod tests { | 3 | 1 | 9 | 20 | 2 | 80 | +----+----+----+----+----+----+ "); + + assert_join_metrics!(metrics, 6); Ok(()) } @@ -884,7 +899,7 @@ mod tests { Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, ); - let (_, batches) = + let (_, batches, metrics) = join_collect(left, right, on, Operator::GtEq, JoinType::Inner).await?; assert_snapshot!(batches_to_string(&batches), @r" @@ -901,6 +916,8 @@ mod tests { | 3 | 4 | 9 | 10 | 3 | 70 | +----+----+----+----+----+----+ "); + + assert_join_metrics!(metrics, 8); Ok(()) } @@ -933,10 +950,12 @@ mod tests { Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, ); - let (_, batches) = + let (_, batches, metrics) = join_collect(left, right, on, Operator::LtEq, JoinType::Inner).await?; // An empty join result produces no batches at all, not an empty batch. assert!(batches.is_empty()); + + assert_join_metrics!(metrics, 0); Ok(()) } @@ -971,7 +990,7 @@ mod tests { Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, ); - let (_, batches) = + let (_, batches, metrics) = join_collect(left, right, on, Operator::GtEq, JoinType::Full).await?; assert_snapshot!(batches_to_string(&batches), @r" @@ -984,6 +1003,7 @@ mod tests { +----+----+-----+----+----+-----+ "); + assert_join_metrics!(metrics, 3); Ok(()) } @@ -1020,7 +1040,7 @@ mod tests { Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, ); - let (_, batches) = + let (_, batches, metrics) = join_collect(left, right, on, Operator::Gt, JoinType::Left).await?; assert_snapshot!(batches_to_string(&batches), @r" @@ -1035,6 +1055,8 @@ mod tests { | 1 | 1 | 7 | | | | +----+----+----+----+----+----+ "); + + assert_join_metrics!(metrics, 6); Ok(()) } @@ -1071,7 +1093,7 @@ mod tests { Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, ); - let (_, batches) = + let (_, batches, metrics) = join_collect(left, right, on, Operator::Gt, JoinType::Right).await?; assert_snapshot!(batches_to_string(&batches), @r" @@ -1084,6 +1106,8 @@ mod tests { | | | | 10 | 5 | 70 | +----+----+----+----+----+----+ "); + + assert_join_metrics!(metrics, 4); Ok(()) } @@ -1120,7 +1144,7 @@ mod tests { Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, ); - let (_, batches) = + let (_, batches, metrics) = join_collect(left, right, on, Operator::Lt, JoinType::Right).await?; assert_snapshot!(batches_to_string(&batches), @r" @@ -1134,6 +1158,8 @@ mod tests { | 3 | 1 | 9 | 10 | 2 | 70 | +----+----+----+----+----+----+ "); + + assert_join_metrics!(metrics, 5); Ok(()) } @@ -1170,7 +1196,7 @@ mod tests { Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, ); - let (_, batches) = + let (_, batches, metrics) = join_collect(left, right, on, Operator::LtEq, JoinType::Inner).await?; // Expected grouping follows right.b1 descending (4, 3, 2) @@ -1185,6 +1211,8 @@ mod tests { | 3 | 2 | 9 | 30 | 2 | 90 | +----+----+----+----+----+----+ "); + + assert_join_metrics!(metrics, 5); Ok(()) } @@ -1221,7 +1249,7 @@ mod tests { Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, ); - let (_, batches) = + let (_, batches, metrics) = join_collect(left, right, on, Operator::Gt, JoinType::Inner).await?; // Grouped by right in ascending evaluation for > (1,2,3) @@ -1235,6 +1263,8 @@ mod tests { | 3 | 4 | 9 | 10 | 3 | 70 | +----+----+----+----+----+----+ "); + + assert_join_metrics!(metrics, 4); Ok(()) } @@ -1334,6 +1364,16 @@ mod tests { | | | | 20 | 5 | 80 | +----+----+----+----+----+----+ "); + + let metrics = join.metrics().expect("metrics should be available"); + assert_join_metrics!(metrics, 2); + assert!( + metrics + .sum_by_name("join_time") + .expect("join_time metric") + .as_usize() + > 0 + ); Ok(()) } @@ -1352,7 +1392,7 @@ mod tests { Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, ); - let (_, batches) = + let (_, batches, metrics) = join_collect(left, right, on, Operator::Lt, JoinType::Inner).await?; assert_snapshot!(batches_to_string(&batches), @r" @@ -1362,6 +1402,8 @@ mod tests { | 42 | 5 | 999 | 30 | 7 | 90 | +----+----+-----+----+----+----+ "); + + assert_join_metrics!(metrics, 1); Ok(()) } @@ -1384,11 +1426,13 @@ mod tests { Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, ); - let (_, batches) = + let (_, batches, metrics) = join_collect(left, right, on, Operator::Gt, JoinType::Inner).await?; // An empty join result produces no batches at all, not an empty batch. assert!(batches.is_empty()); + + assert_join_metrics!(metrics, 0); Ok(()) } @@ -1425,7 +1469,7 @@ mod tests { Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, ); - let (_, batches) = + let (_, batches, metrics) = join_collect(left, right, on, Operator::Lt, JoinType::Inner).await?; assert_snapshot!(batches_to_string(&batches), @r" @@ -1435,6 +1479,8 @@ mod tests { | 1970-01-04 | 2022-04-23 | 1970-01-10 | 1970-01-31 | 2022-04-25 | 1970-04-01 | +------------+------------+------------+------------+------------+------------+ "); + + assert_join_metrics!(metrics, 1); Ok(()) } @@ -1471,7 +1517,7 @@ mod tests { Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, ); - let (_, batches) = + let (_, batches, metrics) = join_collect(left, right, on, Operator::Lt, JoinType::Inner).await?; assert_snapshot!(batches_to_string(&batches), @r" @@ -1481,6 +1527,8 @@ mod tests { | 1970-01-01T00:00:00.003 | 2022-04-23T08:44:01 | 1970-01-01T00:00:00.009 | 1970-01-01T00:00:00.030 | 2022-04-25T16:17:21 | 1970-01-01T00:00:00.090 | +-------------------------+---------------------+-------------------------+-------------------------+---------------------+-------------------------+ "); + + assert_join_metrics!(metrics, 1); Ok(()) } @@ -1515,7 +1563,7 @@ mod tests { Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, ); - let (_, batches) = + let (_, batches, metrics) = join_collect(left, right, on, Operator::Lt, JoinType::Right).await?; assert_snapshot!(batches_to_string(&batches), @r" @@ -1526,6 +1574,159 @@ mod tests { | | | | 1970-01-01T00:00:00.010 | 2022-04-23T08:44:01 | 1970-01-01T00:00:00.080 | +-------------------------+---------------------+-------------------------+-------------------------+---------------------+-------------------------+ "); + + assert_join_metrics!(metrics, 2); Ok(()) } + + /// Wraps an input plan and sleeps `delay` before yielding each of its batches, to + /// simulate a slow streamed input. + #[derive(Debug)] + struct DelayedExec { + input: Arc, + delay: Duration, + } + + impl DisplayAs for DelayedExec { + fn fmt_as( + &self, + _t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "DelayedExec") + } + } + + impl ExecutionPlan for DelayedExec { + fn name(&self) -> &str { + "DelayedExec" + } + + fn properties(&self) -> &Arc { + self.input.properties() + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + Ok(Arc::new(Self { + input: Arc::clone(&children[0]), + delay: self.delay, + })) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let stream = self.input.execute(partition, context)?; + let schema = stream.schema(); + let delay = self.delay; + Ok(Box::pin(RecordBatchStreamAdapter::new( + schema, + stream.then(move |item| async move { + tokio::time::sleep(delay).await; + item + }), + ))) + } + } + + fn join_time_of(metrics: &MetricsSet) -> Duration { + Duration::from_nanos( + metrics + .sum_by_name("join_time") + .map(|m| m.as_usize()) + .unwrap_or(0) as u64, + ) + } + + /// `join_time` must not include time spent waiting for the streamed input: the timer + /// in `fetch_stream_batch` starts only once the input's poll returns `Ready`. + /// + /// Retries with 4x the delay (up to 3 attempts) when the `join_time < delay` check + /// fails. This de-flakes the check without masking real bugs: a genuine exclusion bug + /// makes `join_time` absorb the injected waits, so it scales with the delay and fails + /// at every escalation level, while a fixed-size disturbance (e.g. the OS preempting + /// the thread while a `join_time` timer is running) cannot grow 4x with it. + /// Deterministic invariants (row count, wall-time lower bound) are asserted on every + /// run and never retried. Mirrors the sort-merge join's `check_join_time_excluded`. + #[tokio::test] + async fn join_time_excludes_streamed_input_wait() -> Result<()> { + let mut delay = Duration::from_millis(50); + for attempt in 0..3 { + // Same data as `join_inner_less_than`, with the streamed side split into + // three batches so each one incurs the injected delay. + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![3, 2, 1]), + ("c1", &vec![7, 8, 9]), + ); + + let streamed_schema = Schema::new(vec![ + Field::new("a2", DataType::Int32, false), + Field::new("b1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ]); + let streamed_batches = vec![ + build_table_i32(("a2", &vec![10]), ("b1", &vec![2]), ("c2", &vec![70])), + build_table_i32(("a2", &vec![20]), ("b1", &vec![3]), ("c2", &vec![80])), + build_table_i32(("a2", &vec![30]), ("b1", &vec![4]), ("c2", &vec![90])), + ]; + let right_mem = TestMemoryExec::try_new_exec( + &[streamed_batches], + Arc::new(streamed_schema), + None, + )?; + let right = Arc::new(DelayedExec { + input: right_mem, + delay, + }); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + let join = join(left, right, on, Operator::Lt, JoinType::Inner)?; + + let start = Instant::now(); + let batches = + common::collect(join.execute(0, Arc::new(TaskContext::default()))?) + .await?; + let wall = start.elapsed(); + + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 6, "all streamed rows should find their matches"); + assert!( + wall >= delay * 3, + "streamed delays should dominate wall time, got {wall:?}" + ); + + let join_time = + join_time_of(&join.metrics().expect("metrics should be available")); + if join_time < delay { + return Ok(()); + } + assert!( + attempt < 2, + "join_time ({join_time:?}) should be well below the injected \ + delay ({delay:?}) even after escalating retries; wall {wall:?}" + ); + delay *= 4; + } + unreachable!() + } } diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/existence_join.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/existence_join.rs index b2c5212999f3b..947569203ffca 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/existence_join.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/existence_join.rs @@ -227,7 +227,10 @@ impl ExistencePWMJStream { return Poll::Ready(Ok(StatefulStreamResult::Continue)); } - match ready!(self.streamed.poll_next_unpin(cx)) { + let next_batch = ready!(self.streamed.poll_next_unpin(cx)); + let join_time = self.join_metrics.join_time.clone(); + let _join_timer = join_time.timer(); + match next_batch { None => self.finish_streamed_side()?, Some(Ok(batch)) => { let stream_values: ArrayRef = self @@ -389,6 +392,7 @@ impl ExistencePWMJStream { /// Emits the existence result by slicing at the watermark: the marked buffered rows for /// `LeftSemi`, the unmarked ones for `LeftAnti`. fn emit_matched(&mut self) -> Result>> { + let _join_timer = self.join_metrics.join_time.timer(); if !self.emitted { self.emitted = true; @@ -485,7 +489,7 @@ mod tests { use crate::{ ExecutionPlan, common, joins::PiecewiseMergeJoinExec, - test::{TestMemoryExec, build_table_i32}, + test::{TestMemoryExec, assert_join_metrics, build_table_i32}, }; use arrow_schema::{DataType, Field, Schema}; use datafusion_common::test_util::batches_to_string; @@ -638,7 +642,8 @@ mod tests { /// The final pass pushes one slice of the buffered batch into a `BatchCoalescer`, so a /// result wider than `batch_size` has to be drained over several polls. Also pins the - /// `output_rows` metric, which `record_poll` in `poll_next` is what supplies. + /// `output_rows` metric, which `record_poll` in `poll_next` is what supplies, and + /// `join_time`, which the streamed-side scan and the final pass record. #[tokio::test] async fn final_pass_chunks_output_and_records_output_rows() -> Result<()> { let left = build_table( @@ -685,12 +690,15 @@ mod tests { assert_eq!(batches.len(), 2, "expected the final pass to be chunked"); assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 4); - let output_rows = join - .metrics() - .unwrap() - .output_rows() - .expect("output_rows metric"); - assert_eq!(output_rows, 4); + let metrics = join.metrics().expect("metrics should be available"); + assert_join_metrics!(metrics, 4); + assert!( + metrics + .sum_by_name("join_time") + .expect("join_time metric") + .as_usize() + > 0 + ); Ok(()) } diff --git a/datafusion/sqllogictest/test_files/explain_analyze.slt b/datafusion/sqllogictest/test_files/explain_analyze.slt index e083e8be33e95..81b274d6c1377 100644 --- a/datafusion/sqllogictest/test_files/explain_analyze.slt +++ b/datafusion/sqllogictest/test_files/explain_analyze.slt @@ -422,7 +422,7 @@ FROM t1 JOIN t2 ON t1.k > t2.k; ---- Plan with Metrics -01)PiecewiseMergeJoin: operator=Gt, join_type=Inner, on=(k > k), metrics=[output_bytes=0.0 B, build_mem_used=144.0 B] +01)PiecewiseMergeJoin: operator=Gt, join_type=Inner, on=(k > k), metrics=[output_bytes=128.0 KB, build_mem_used=144.0 B] 02)--ProjectionExec: expr=[column1@0 as k], metrics=[output_bytes=16.0 B] 03)----SortExec: expr=[column1@0 ASC], preserve_partitioning=[false], metrics=[output_bytes=16.0 B, spilled_bytes=0.0 B] 04)------DataSourceExec: partitions=1, partition_sizes=[1], metrics=[]