From aa2d37b8ed61833e3e9ebf7b9817b0c1d31bb068 Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Mon, 31 Aug 2026 13:38:18 +0530 Subject: [PATCH 1/2] fix(pwmj): record output/join_time/probe metrics --- .../piecewise_merge_join/classic_join.rs | 156 +++++++++++++++- .../piecewise_merge_join/existence_join.rs | 169 +++++++++++++++++- 2 files changed, 318 insertions(+), 7 deletions(-) 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 ac1385223f0fa..cd5d0741b3f87 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 @@ -240,6 +240,7 @@ impl ClassicPWMJStream { self.join_metrics.input_batches.add(1); self.join_metrics.input_rows.add(batch.num_rows()); + self.join_metrics.probe_hit_rate.add_total(batch.num_rows()); // Sort stream values and change the streamed record batch accordingly let indices = sort_to_indices( @@ -291,6 +292,7 @@ impl ClassicPWMJStream { } // Produce more work + let join_timer = self.join_metrics.join_time.timer(); let batch = resolve_classic_join( buffered_side, stream_batch, @@ -299,7 +301,9 @@ impl ClassicPWMJStream { self.sort_option, self.join_type, &mut self.batch_process_state, + &self.join_metrics, )?; + join_timer.done(); if !self.batch_process_state.continue_process { // Scan finished; re-enter through the drain guard above. @@ -332,6 +336,7 @@ impl ClassicPWMJStream { let buffered_data = Arc::clone(&self.buffered_side.try_as_ready()?.buffered_data); let buffered_batch = buffered_data.batch(); + let join_timer = self.join_metrics.join_time.timer(); // Every match marks the suffix `[k, buffered_len)`, so the buffered rows that were // never matched are exactly the complementary prefix `[0, min_marked)` -- which // includes the null-keyed rows, since nulls sort first and the scan starts past @@ -354,6 +359,7 @@ impl ClassicPWMJStream { buffered_columns.extend(streamed_columns); let batch = RecordBatch::try_new(Arc::clone(&self.schema), buffered_columns)?; + join_timer.done(); self.batch_process_state.output_batches.push_batch(batch)?; @@ -425,11 +431,15 @@ impl Stream for ClassicPWMJStream { mut self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>, ) -> Poll> { - self.poll_next_impl(cx) + // `record_poll` fills in `output_rows` and `end_time`; `elapsed_compute` is handled + // by `BuildProbeJoinMetrics::drop`. + let poll = self.poll_next_impl(cx); + self.join_metrics.baseline.record_poll(poll) } } // For Left, Right, Full, and Inner joins, incoming stream batches will already be sorted. +#[expect(clippy::too_many_arguments)] fn resolve_classic_join( buffered_side: &mut BufferedSideReadyState, stream_batch: &SortedStreamBatch, @@ -438,6 +448,7 @@ fn resolve_classic_join( sort_options: SortOptions, join_type: JoinType, batch_process_state: &mut BatchProcessState, + join_metrics: &BuildProbeJoinMetrics, ) -> Result { let buffered_len = buffered_side.buffered_data.values().len(); let stream_values = stream_batch.compare_key_values(); @@ -483,6 +494,9 @@ fn resolve_classic_join( if compare == Ordering::Less { batch_process_state.found = true; let count = buffered_len - buffer_idx; + join_metrics.probe_hit_rate.add_part(1); + join_metrics.avg_fanout.add_part(count); + join_metrics.avg_fanout.add_total(1); let batch = build_matched_indices_and_mark_buffered( (buffer_idx, count), @@ -513,6 +527,9 @@ fn resolve_classic_join( if matches!(compare, Ordering::Equal | Ordering::Less) { batch_process_state.found = true; let count = buffered_len - buffer_idx; + join_metrics.probe_hit_rate.add_part(1); + join_metrics.avg_fanout.add_part(count); + join_metrics.avg_fanout.add_total(1); let batch = build_matched_indices_and_mark_buffered( (buffer_idx, count), (row_idx, count), @@ -1550,4 +1567,141 @@ mod tests { "); Ok(()) } + + fn ratio_metric(metrics: &crate::metrics::MetricsSet, name: &str) -> (usize, usize) { + metrics + .iter() + .find_map(|m| match m.value() { + crate::metrics::MetricValue::Ratio { + name: metric_name, + ratio_metrics, + } if metric_name == name => { + Some((ratio_metrics.part(), ratio_metrics.total())) + } + _ => None, + }) + .unwrap_or_else(|| panic!("{name} metric not found")) + } + + fn sum_metric( + metrics: &crate::metrics::MetricsSet, + matches: impl Fn(&crate::metrics::MetricValue) -> Option, + ) -> usize { + metrics.iter().filter_map(|m| matches(m.value())).sum() + } + + /// Classic joins never routed `poll_next` through `record_poll`, so `output_rows`, + /// `output_bytes` and `output_batches` stayed at zero regardless of how many rows the + /// join actually produced. Also pins `probe_hit_rate`/`avg_fanout`, which classic join + /// never populated at all. + #[tokio::test] + async fn inner_join_records_output_and_probe_metrics() -> Result<()> { + // Buffered side must already be ascending for `Gt`, so this is the one classic-join + // test that hand-derives expected counts rather than only checking output rows. + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![1, 2, 5]), + ("c1", &vec![7, 8, 9]), + ); + // Fed unsorted; `fetch_stream_batch` sorts each streamed batch internally. + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 3, 2]), + ("c2", &vec![70, 80, 90]), + ); + + 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::Gt, JoinType::Inner)?; + let stream = join.execute(0, Arc::new(TaskContext::default()))?; + let batches = common::collect(stream).await?; + + // Every streamed value (2, 3, 4) is only exceeded by buffered value 5, so each + // produces exactly one matched row. + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 3); + + let metrics = join.metrics().unwrap(); + assert_eq!( + metrics.output_rows(), + Some(3), + "output_rows must reflect the batches actually produced" + ); + assert!( + sum_metric(&metrics, |v| match v { + crate::metrics::MetricValue::OutputBytes(c) => Some(c.value()), + _ => None, + }) > 0 + ); + assert_eq!( + sum_metric(&metrics, |v| match v { + crate::metrics::MetricValue::OutputBatches(c) => Some(c.value()), + _ => None, + }), + 1 + ); + + // All 3 streamed rows found a match. + assert_eq!(ratio_metric(&metrics, "probe_hit_rate"), (3, 3)); + // Each match was against exactly one buffered row (value 5). + assert_eq!(ratio_metric(&metrics, "avg_fanout"), (3, 3)); + + Ok(()) + } + + /// `Full` join is the only join type that runs both `join_time`-wrapped code paths: + /// `resolve_classic_join`'s matched/unmatched-streamed pass, and + /// `process_unmatched_buffered_batch`'s unmatched-buffered pass. `probe_hit_rate` and + /// `avg_fanout` must count only the 3 real matches -- the 2 unmatched buffered rows the + /// second pass adds to the output must not leak into either ratio. + #[tokio::test] + async fn full_join_unmatched_buffered_rows_do_not_pollute_probe_metrics() -> Result<()> + { + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![1, 2, 5]), + ("c1", &vec![7, 8, 9]), + ); + // Fed unsorted; `fetch_stream_batch` sorts each streamed batch internally. Every + // value here is less than the buffered maximum (5), so every streamed row matches + // and buffered rows 1 and 2 are left unmatched. + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 3, 2]), + ("c2", &vec![70, 80, 90]), + ); + + 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::Gt, JoinType::Full)?; + let stream = join.execute(0, Arc::new(TaskContext::default()))?; + let batches = common::collect(stream).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | 3 | 5 | 9 | 30 | 2 | 90 | + | 3 | 5 | 9 | 20 | 3 | 80 | + | 3 | 5 | 9 | 10 | 4 | 70 | + | 1 | 1 | 7 | | | | + | 2 | 2 | 8 | | | | + +----+----+----+----+----+----+ + "); + + let metrics = join.metrics().unwrap(); + // 3 matched rows + 2 unmatched-buffered rows. + assert_eq!(metrics.output_rows(), Some(5)); + + // Still only the 3 real matches -- the unmatched-buffered pass leaves these alone. + assert_eq!(ratio_metric(&metrics, "probe_hit_rate"), (3, 3)); + assert_eq!(ratio_metric(&metrics, "avg_fanout"), (3, 3)); + + Ok(()) + } } 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 887367b22eeb2..321fc0ae89c56 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 @@ -238,12 +238,19 @@ impl ExistencePWMJStream { self.join_metrics.input_batches.add(1); self.join_metrics.input_rows.add(batch.num_rows()); - // Only the batch's extreme key is ever compared against the buffered side, - // so reduce the batch to that one key. - let stream_values = - extreme_key(&stream_values, self.sort_option.descending)?; - - self.mark_matched_buffered_rows(&stream_values)?; + // An empty batch has no extreme key to compare, so it can neither match + // nor miss -- counting it either way would understate the real hit rate. + if batch.num_rows() > 0 { + let join_time = self.join_metrics.join_time.clone(); + let join_timer = join_time.timer(); + // Only the batch's extreme key is ever compared against the buffered + // side, so reduce the batch to that one key. + let stream_values = + extreme_key(&stream_values, self.sort_option.descending)?; + + self.mark_matched_buffered_rows(&stream_values)?; + join_timer.done(); + } } Some(Err(err)) => return Poll::Ready(Err(err)), } @@ -291,6 +298,7 @@ impl ExistencePWMJStream { fn mark_matched_buffered_rows(&mut self, stream_values: &ArrayRef) -> Result<()> { let operator = self.operator; let sort_option = self.sort_option; + self.join_metrics.probe_hit_rate.add_total(1); { let buffered_data = &self.buffered_side.try_as_ready()?.buffered_data; @@ -362,6 +370,7 @@ impl ExistencePWMJStream { // batch matches nothing new. let buffer_idx = lo; if buffer_idx < scan_limit { + self.join_metrics.probe_hit_rate.add_part(1); // Everything from `buffer_idx` on matches, so lowering the // watermark to it records the match: the marked set is exactly // `[min_marked, buffered_len)` and needs no bitmap. @@ -814,4 +823,152 @@ mod tests { } Ok(()) } + + /// Existence join never populated `probe_hit_rate`, so a streamed batch whose extreme + /// key failed to lower the watermark was indistinguishable from one that did. Two + /// batches here: the first lowers the watermark, the second lands entirely inside the + /// already-marked region and must count as a miss. + #[tokio::test] + async fn probe_hit_rate_counts_batches_that_advance_the_watermark() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3, 4, 5]), + ("b1", &vec![1, 2, 3, 4, 5]), + ("c1", &vec![10, 20, 30, 40, 50]), + ); + + let streamed_schema = Schema::new(vec![ + Field::new("a2", DataType::Int32, false), + Field::new("b1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ]); + // b1=3 lowers the watermark to buffered index 3 (value 4). + let batch1 = + build_table_i32(("a2", &vec![10]), ("b1", &vec![3]), ("c2", &vec![70])); + // b1=4 only matches within the already-marked suffix, so it can't lower the + // watermark further -- a miss. + let batch2 = + build_table_i32(("a2", &vec![20]), ("b1", &vec![4]), ("c2", &vec![80])); + let right = TestMemoryExec::try_new_exec( + &[vec![batch1, batch2]], + Arc::new(streamed_schema), + None, + )?; + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + let join = PiecewiseMergeJoinExec::try_new( + left, + right, + on, + Operator::Gt, + JoinType::LeftSemi, + 1, + )?; + + let stream = join.execute(0, Arc::new(TaskContext::default()))?; + let batches = common::collect(stream).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 4 | 4 | 40 | + | 5 | 5 | 50 | + +----+----+----+ + "); + + let metrics = join.metrics().unwrap(); + let hit_rate = metrics + .iter() + .find_map(|m| match m.value() { + crate::metrics::MetricValue::Ratio { + name, + ratio_metrics, + } if name == "probe_hit_rate" => { + Some((ratio_metrics.part(), ratio_metrics.total())) + } + _ => None, + }) + .expect("probe_hit_rate metric"); + assert_eq!(hit_rate, (1, 2), "one hit, one miss"); + + Ok(()) + } + + /// An empty streamed batch has no extreme key to compare against the buffered side, + /// so it must count as neither a hit nor a miss. Before the guard in `scan_stream_batch`, + /// `mark_matched_buffered_rows` ran unconditionally and inflated `probe_hit_rate`'s + /// denominator with misses for batches that never actually scanned anything. + #[tokio::test] + async fn probe_hit_rate_ignores_empty_streamed_batches() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![1, 2, 5]), + ("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 empty_batch = build_table_i32( + ("a2", &Vec::new()), + ("b1", &Vec::new()), + ("c2", &Vec::new()), + ); + let real_batch = + build_table_i32(("a2", &vec![10]), ("b1", &vec![0]), ("c2", &vec![70])); + let right = TestMemoryExec::try_new_exec( + &[vec![empty_batch, real_batch]], + Arc::new(streamed_schema), + None, + )?; + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + let join = PiecewiseMergeJoinExec::try_new( + left, + right, + on, + Operator::Gt, + JoinType::LeftSemi, + 1, + )?; + + let stream = join.execute(0, Arc::new(TaskContext::default()))?; + let batches = common::collect(stream).await?; + + // b1=0 is below every buffered value, so all three buffered rows match. + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 1 | 1 | 7 | + | 2 | 2 | 8 | + | 3 | 5 | 9 | + +----+----+----+ + "); + + let metrics = join.metrics().unwrap(); + let hit_rate = metrics + .iter() + .find_map(|m| match m.value() { + crate::metrics::MetricValue::Ratio { + name, + ratio_metrics, + } if name == "probe_hit_rate" => { + Some((ratio_metrics.part(), ratio_metrics.total())) + } + _ => None, + }) + .expect("probe_hit_rate metric"); + // Only the real batch counts -- the empty batch contributes to neither part nor total. + assert_eq!(hit_rate, (1, 1)); + + Ok(()) + } } From a0ba9a7897a11f178e009921a5fbb101492ab061 Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Sun, 6 Sep 2026 10:54:54 +0530 Subject: [PATCH 2/2] remove nested join timer --- .../src/joins/piecewise_merge_join/classic_join.rs | 4 ---- .../src/joins/piecewise_merge_join/existence_join.rs | 3 --- 2 files changed, 7 deletions(-) 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 96548af1e4210..dc70040928d6d 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 @@ -295,7 +295,6 @@ impl ClassicPWMJStream { } // Produce more work - let join_timer = self.join_metrics.join_time.timer(); let batch = resolve_classic_join( buffered_side, stream_batch, @@ -306,7 +305,6 @@ impl ClassicPWMJStream { &mut self.batch_process_state, &self.join_metrics, )?; - join_timer.done(); if !self.batch_process_state.continue_process { // Scan finished; re-enter through the drain guard above. @@ -340,7 +338,6 @@ impl ClassicPWMJStream { let buffered_data = Arc::clone(&self.buffered_side.try_as_ready()?.buffered_data); let buffered_batch = buffered_data.batch(); - let join_timer = self.join_metrics.join_time.timer(); // Every match marks the suffix `[k, buffered_len)`, so the buffered rows that were // never matched are exactly the complementary prefix `[0, min_marked)` -- which // includes the null-keyed rows, since nulls sort first and the scan starts past @@ -363,7 +360,6 @@ impl ClassicPWMJStream { buffered_columns.extend(streamed_columns); let batch = RecordBatch::try_new(Arc::clone(&self.schema), buffered_columns)?; - join_timer.done(); self.batch_process_state.output_batches.push_batch(batch)?; 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 b077a92157c55..15424caa3201c 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 @@ -244,15 +244,12 @@ impl ExistencePWMJStream { // An empty batch has no extreme key to compare, so it can neither match // nor miss -- counting it either way would understate the real hit rate. if batch.num_rows() > 0 { - let join_time = self.join_metrics.join_time.clone(); - let join_timer = join_time.timer(); // Only the batch's extreme key is ever compared against the buffered // side, so reduce the batch to that one key. let stream_values = extreme_key(&stream_values, self.sort_option.descending)?; self.mark_matched_buffered_rows(&stream_values)?; - join_timer.done(); } } Some(Err(err)) => return Poll::Ready(Err(err)),