From 3cfbb98737e3d238485886dc705098bd395fe790 Mon Sep 17 00:00:00 2001 From: Yu-Chuan Hung Date: Sun, 13 Sep 2026 22:59:24 +0800 Subject: [PATCH] fix: count each probe row once in HashJoinExec probe_hit_rate and avg_fanout - `probe_hit_rate` added the batch's row count to its total on every chunk. - A probe row whose matches span a chunk boundary was counted in both chunks, inflating the `probe_hit_rate` part and the `avg_fanout` total. `ProcessProbeBatchState` now tracks the last counted probe index. --- .../physical-plan/src/joins/hash_join/exec.rs | 142 ++++++++++++++++++ .../src/joins/hash_join/stream.rs | 41 ++++- 2 files changed, 175 insertions(+), 8 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index b72e180543f9a..eceb06a4cf403 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -3027,6 +3027,25 @@ mod tests { } } + #[track_caller] + fn assert_ratio_metric( + metrics: &MetricsSet, + metric_name: &str, + expected_part: usize, + expected_total: usize, + ) { + let Some(MetricValue::Ratio { ratio_metrics, .. }) = + metrics.sum_by_name(metric_name) + else { + panic!("should have {metric_name} metrics") + }; + assert_eq!( + (ratio_metrics.part(), ratio_metrics.total()), + (expected_part, expected_total), + "{metric_name} (part, total) mismatch", + ); + } + fn build_schema_and_on() -> Result<(SchemaRef, SchemaRef, JoinOn)> { let left_schema = Arc::new(Schema::new(vec![ Field::new("a1", DataType::Int32, true), @@ -3073,6 +3092,7 @@ mod tests { use datafusion_physical_expr::{ EquivalenceProperties, PhysicalSortExpr, RangePartitioning, SplitPoint, }; + use datafusion_physical_expr_common::metrics::MetricValue; use futures::StreamExt; use hashbrown::HashTable; use insta::{allow_duplicates, assert_snapshot}; @@ -4079,6 +4099,128 @@ mod tests { TestMemoryExec::try_new_exec(&[vec![batch.clone(), batch]], schema, None).unwrap() } + /// `probe_hit_rate` and `avg_fanout` must count each probe row once, even + /// when a probe batch is processed in several chunks. Every probe row matches + /// all 3 build rows, so any `batch_size` below 9 splits the probe batch into + /// chunks, and some splits cut a single row's matches across chunks. + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_probe_metrics_count_each_probe_row_once( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![2, 2, 2]), + ("c1", &vec![3, 4, 5]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![2, 2, 2]), + ("c2", &vec![30, 40, 50]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (columns, batches, metrics) = join_collect( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + &JoinType::Inner, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b1", "c2"]); + allow_duplicates! { + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | 1 | 2 | 3 | 10 | 2 | 30 | + | 2 | 2 | 4 | 10 | 2 | 30 | + | 3 | 2 | 5 | 10 | 2 | 30 | + | 1 | 2 | 3 | 20 | 2 | 40 | + | 2 | 2 | 4 | 20 | 2 | 40 | + | 3 | 2 | 5 | 20 | 2 | 40 | + | 1 | 2 | 3 | 30 | 2 | 50 | + | 2 | 2 | 4 | 30 | 2 | 50 | + | 3 | 2 | 5 | 30 | 2 | 50 | + +----+----+----+----+----+----+ + "); + } + + assert_join_metrics!(metrics, 9); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + assert_ratio_metric(&metrics, "probe_hit_rate", 3, 3); + assert_ratio_metric(&metrics, "avg_fanout", 9, 3); + + Ok(()) + } + + /// Complements `join_probe_metrics_count_each_probe_row_once`: with unique + /// build keys each probe row has at most one match, so chunks always split + /// between probe rows. A probe row that starts a new chunk must still be + /// counted, even though the lookup offset already points at it. + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_probe_metrics_count_probe_row_starting_new_chunk( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 6]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 25, 30]), + ("b1", &vec![4, 4, 4, 40]), + ("c2", &vec![70, 80, 85, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (columns, batches, metrics) = join_collect( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + &JoinType::Inner, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b1", "c2"]); + allow_duplicates! { + assert_snapshot!(batches_to_string(&batches), @" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | 1 | 4 | 7 | 10 | 4 | 70 | + | 1 | 4 | 7 | 20 | 4 | 80 | + | 1 | 4 | 7 | 25 | 4 | 85 | + +----+----+----+----+----+----+ + "); + } + + assert_join_metrics!(metrics, 3); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + assert_ratio_metric(&metrics, "probe_hit_rate", 3, 4); + assert_ratio_metric(&metrics, "avg_fanout", 3, 3); + + Ok(()) + } + #[apply(hash_join_exec_configs)] #[tokio::test] async fn join_left_multi_batch( diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs index 625661188d81c..7b70b2bf3b318 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -188,6 +188,11 @@ pub(super) struct ProcessProbeBatchState { offset: MapOffset, /// Max joined probe-side index from current batch joined_probe_idx: Option, + /// Max probe-side index with a join-key match from current batch, before + /// the join filter is applied (unlike `joined_probe_idx`). Lets + /// `probe_hit_rate` and `avg_fanout` count a probe row whose matches span + /// several chunks only once. + matched_probe_idx: Option, } impl ProcessProbeBatchState { @@ -197,6 +202,23 @@ impl ProcessProbeBatchState { self.joined_probe_idx = joined_probe_idx; } } + + /// Returns how many probe rows in `right_indices`, the join-key matches of + /// the current chunk, have not been counted by a previous chunk of this + /// probe batch, and records the last matched index in `matched_probe_idx`. + fn count_new_matched_probe_rows(&mut self, right_indices: &UInt32Array) -> usize { + let values = right_indices.values(); + let mut count = count_distinct_sorted_indices(right_indices); + // A probe row whose matches span a chunk boundary is the last index of the + // previous chunk and the first index of this one; count it only once. + if let (Some(&first), Some(&last)) = (values.first(), values.last()) { + if Some(first) == self.matched_probe_idx { + count -= 1; + } + self.matched_probe_idx = Some(last); + } + count + } } /// Container for HashJoinStreamState::EmitUnmatchedBuildRows related data @@ -786,6 +808,7 @@ impl HashJoinStream { valid_keys, offset: (0, None), joined_probe_idx: None, + matched_probe_idx: None, }); } Some(Err(err)) => return Poll::Ready(Err(err)), @@ -803,9 +826,13 @@ impl HashJoinStream { let state = self.state.try_as_process_probe_batch_mut()?; let build_side = self.build_side.try_as_ready_mut()?; - self.join_metrics - .probe_hit_rate - .add_total(state.batch.num_rows()); + // A probe batch may be processed in several chunks; count its rows once, + // on the first chunked lookup (offset == (0, None)). + if state.offset == (0, None) { + self.join_metrics + .probe_hit_rate + .add_total(state.batch.num_rows()); + } let timer = self.join_metrics.join_time.timer(); @@ -890,17 +917,15 @@ impl HashJoinStream { } }; - let distinct_right_indices_count = count_distinct_sorted_indices(&right_indices); + let matched_probe_rows = state.count_new_matched_probe_rows(&right_indices); self.join_metrics .probe_hit_rate - .add_part(distinct_right_indices_count); + .add_part(matched_probe_rows); self.join_metrics.avg_fanout.add_part(left_indices.len()); - self.join_metrics - .avg_fanout - .add_total(distinct_right_indices_count); + self.join_metrics.avg_fanout.add_total(matched_probe_rows); // apply join filter if exists let (left_indices, right_indices) = if let Some(filter) = &self.filter {