@@ -706,7 +706,9 @@ impl CrossJoinStream {
706706mod 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