Pwmj metrics accounting

datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs

@@ -242,6 +242,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(
@@ -302,6 +303,7 @@ impl ClassicPWMJStream {
             self.sort_option,
             self.join_type,
             &mut self.batch_process_state,
+            &self.join_metrics,
         )?;
 
         if !self.batch_process_state.continue_process {
@@ -429,12 +431,15 @@ impl Stream for ClassicPWMJStream {
         mut self: std::pin::Pin<&mut Self>,
         cx: &mut std::task::Context<'_>,
     ) -> Poll<Option<Self::Item>> {
+        // `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,
@@ -443,6 +448,7 @@ fn resolve_classic_join(
     sort_options: SortOptions,
     join_type: JoinType,
     batch_process_state: &mut BatchProcessState,
+    join_metrics: &BuildProbeJoinMetrics,
 ) -> Result<RecordBatch> {
     let buffered_len = buffered_side.buffered_data.values().len();
     let stream_values = stream_batch.compare_key_values();
@@ -488,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),
@@ -518,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),
@@ -1601,6 +1613,143 @@ mod tests {
         Ok(())
     }
 
+    fn ratio_metric(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: &MetricsSet,
+        matches: impl Fn(&crate::metrics::MetricValue) -> Option<usize>,
+    ) -> 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::<usize>(), 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(())
+    }
+
     /// Wraps an input plan and sleeps `delay` before yielding each of its batches, to
     /// simulate a slow streamed input.
     #[derive(Debug)]

datafusion/physical-plan/src/joins/piecewise_merge_join/existence_join.rs

@@ -241,12 +241,16 @@ 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 {
+                    // 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)?;
+                }
             }
             Some(Err(err)) => return Poll::Ready(Err(err)),
         }
@@ -294,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;
@@ -365,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.
@@ -822,4 +828,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(())
+    }
 }