Skip to content

Commit e12556c

Browse files
committed
test: cover CrossJoinExec row metrics across batched inputs
Add a Rust test that asserts build_input_rows, input_rows, and output_rows stay stable when CrossJoinExec inputs are split into multiple batches. Related to #25333.
1 parent 6ffc45d commit e12556c

1 file changed

Lines changed: 60 additions & 1 deletion

File tree

‎datafusion/physical-plan/src/joins/cross_join.rs‎

Lines changed: 60 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -706,7 +706,9 @@ impl CrossJoinStream {
706706
mod tests {
707707
use super::*;
708708
use crate::common;
709-
use crate::test::{assert_join_metrics, build_table_scan_i32};
709+
use crate::test::{
710+
TestMemoryExec, assert_join_metrics, build_table_i32, build_table_scan_i32,
711+
};
710712

711713
use datafusion_common::{assert_contains, test_util::batches_to_sort_string};
712714
use datafusion_execution::runtime_env::RuntimeEnvBuilder;
@@ -727,6 +729,29 @@ mod tests {
727729
Ok((columns_header, batches, metrics))
728730
}
729731

732+
/// Split a 3-column i32 table into `RecordBatch`es of at most `input_batch_size` rows.
733+
fn build_table_scan_i32_chunked(
734+
a: (&str, &Vec<i32>),
735+
b: (&str, &Vec<i32>),
736+
c: (&str, &Vec<i32>),
737+
input_batch_size: usize,
738+
) -> Arc<dyn ExecutionPlan> {
739+
let full = build_table_i32(a, b, c);
740+
let n = full.num_rows();
741+
let batches = (0..n)
742+
.step_by(input_batch_size)
743+
.map(|i| full.slice(i, input_batch_size.min(n - i)))
744+
.collect();
745+
TestMemoryExec::try_new_exec(&[batches], full.schema(), None).unwrap()
746+
}
747+
748+
fn metric_count(metrics: &MetricsSet, name: &str) -> usize {
749+
metrics
750+
.sum_by_name(name)
751+
.expect("missing metric")
752+
.as_usize()
753+
}
754+
730755
#[tokio::test]
731756
async fn test_stats_cartesian_product() {
732757
let left_row_count = 11;
@@ -966,6 +991,40 @@ mod tests {
966991
");
967992

968993
assert_join_metrics!(metrics, 6);
994+
assert_eq!(metric_count(&metrics, "build_input_rows"), 3);
995+
assert_eq!(metric_count(&metrics, "input_rows"), 2);
996+
997+
Ok(())
998+
}
999+
1000+
/// `build_input_rows`, `input_rows`, and `output_rows` must count each row
1001+
/// once when the build and probe sides arrive as several batches.
1002+
#[tokio::test]
1003+
async fn join_metrics_count_rows_once_across_batches() -> Result<()> {
1004+
let task_ctx = Arc::new(TaskContext::default());
1005+
// 2: left 2+1, right 2. 1: left 3×1, right 2×1.
1006+
for input_batch_size in [2, 1] {
1007+
let left = build_table_scan_i32_chunked(
1008+
("a1", &vec![1, 2, 3]),
1009+
("b1", &vec![4, 5, 6]),
1010+
("c1", &vec![7, 8, 9]),
1011+
input_batch_size,
1012+
);
1013+
let right = build_table_scan_i32_chunked(
1014+
("a2", &vec![10, 11]),
1015+
("b2", &vec![12, 13]),
1016+
("c2", &vec![14, 15]),
1017+
input_batch_size,
1018+
);
1019+
1020+
let (_, batches, metrics) =
1021+
join_collect(left, right, Arc::clone(&task_ctx)).await?;
1022+
let num_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
1023+
assert_eq!(num_rows, 6);
1024+
assert_join_metrics!(metrics, 6);
1025+
assert_eq!(metric_count(&metrics, "build_input_rows"), 3);
1026+
assert_eq!(metric_count(&metrics, "input_rows"), 2);
1027+
}
9691028

9701029
Ok(())
9711030
}

0 commit comments

Comments
 (0)