@@ -3482,6 +3482,7 @@ pub(crate) mod tests {
34823482 use arrow:: compute:: SortOptions ;
34833483 use arrow:: datatypes:: { DataType , Field } ;
34843484 use datafusion_common:: assert_contains;
3485+ use datafusion_common:: cast:: as_int32_array;
34853486 use datafusion_common:: test_util:: batches_to_sort_string;
34863487 use datafusion_execution:: runtime_env:: RuntimeEnvBuilder ;
34873488 use datafusion_expr:: Operator ;
@@ -3630,6 +3631,176 @@ pub(crate) mod tests {
36303631 assert_eq ! ( left_data. locate( 12 ) , None ) ;
36313632 }
36323633
3634+ const CHUNKED_LEFT_ROWS : i32 = 36 ;
3635+ /// Left rows that [`chunked_right_table`] matches: the first and last row, and the rows on
3636+ /// both sides of a chunk boundary in every layout the test builds (11|12 for 4- and 12-row
3637+ /// chunks, 27|28 for 4- and 7-row chunks). Every other boundary has unmatched rows on both
3638+ /// sides.
3639+ const CHUNKED_LEFT_MATCHED : [ i32 ; 6 ] = [ 0 , 11 , 12 , 27 , 28 , 35 ] ;
3640+
3641+ fn chunked_left_table ( batch_rows : Option < usize > ) -> Arc < dyn ExecutionPlan > {
3642+ let ids: Vec < i32 > = ( 0 ..CHUNKED_LEFT_ROWS ) . collect ( ) ;
3643+ build_table (
3644+ ( "a1" , & ids) ,
3645+ ( "b1" , & ids) ,
3646+ ( "c1" , & ids) ,
3647+ batch_rows,
3648+ Vec :: new ( ) ,
3649+ )
3650+ }
3651+
3652+ /// One row per batch, so a batch size of 12 takes the multi-row probe path (12 / 1 > 10)
3653+ /// and a batch size of 4 the single-row one.
3654+ fn chunked_right_table ( ) -> Arc < dyn ExecutionPlan > {
3655+ let mut ids = CHUNKED_LEFT_MATCHED . to_vec ( ) ;
3656+ ids. push ( 99 ) ;
3657+ build_table (
3658+ ( "a2" , & ids) ,
3659+ ( "b2" , & ids) ,
3660+ ( "c2" , & ids) ,
3661+ Some ( 1 ) ,
3662+ Vec :: new ( ) ,
3663+ )
3664+ }
3665+
3666+ /// left.a1 = right.a2
3667+ fn equality_join_filter ( ) -> JoinFilter {
3668+ let column_indices = vec ! [
3669+ ColumnIndex {
3670+ index: 0 ,
3671+ side: JoinSide :: Left ,
3672+ } ,
3673+ ColumnIndex {
3674+ index: 0 ,
3675+ side: JoinSide :: Right ,
3676+ } ,
3677+ ] ;
3678+ let intermediate_schema = Schema :: new ( vec ! [
3679+ Field :: new( "a1" , DataType :: Int32 , true ) ,
3680+ Field :: new( "a2" , DataType :: Int32 , true ) ,
3681+ ] ) ;
3682+ let expression = Arc :: new ( BinaryExpr :: new (
3683+ Arc :: new ( Column :: new ( "a1" , 0 ) ) ,
3684+ Operator :: Eq ,
3685+ Arc :: new ( Column :: new ( "a2" , 1 ) ) ,
3686+ ) ) as Arc < dyn PhysicalExpr > ;
3687+ JoinFilter :: new ( expression, column_indices, Arc :: new ( intermediate_schema) )
3688+ }
3689+
3690+ /// Number of chunks the build-side load produces for `left` at this target batch size.
3691+ async fn build_side_chunk_count (
3692+ left : & Arc < dyn ExecutionPlan > ,
3693+ target_batch_size : usize ,
3694+ ) -> Result < usize > {
3695+ let task_ctx = Arc :: new ( TaskContext :: default ( ) ) ;
3696+ let reservation = MemoryConsumer :: new ( "test" ) . register ( task_ctx. memory_pool ( ) ) ;
3697+ let metrics = ExecutionPlanMetricsSet :: new ( ) ;
3698+ let load = collect_left_input (
3699+ left. execute ( 0 , task_ctx) ?,
3700+ BuildProbeJoinMetrics :: new ( 0 , & metrics) ,
3701+ reservation,
3702+ false ,
3703+ 1 ,
3704+ None ,
3705+ target_batch_size,
3706+ )
3707+ . await ?;
3708+ let LeftLoad :: InMemory ( data) = load else {
3709+ panic ! ( "the build side fit in memory" ) ;
3710+ } ;
3711+ Ok ( data. chunks . len ( ) )
3712+ }
3713+
3714+ /// The build side arrives either as single-row batches, which the load coalesces into chunks
3715+ /// of exactly `batch_size` rows, or as 7-row batches, which bypass the coalescer and leave
3716+ /// chunk edges off the output batch size, so probe and unmatched-left ranges must be clamped
3717+ /// at them. The same rows delivered as one batch form a single chunk, so the two runs differ
3718+ /// only in chunk layout and must agree.
3719+ #[ rstest]
3720+ #[ tokio:: test]
3721+ async fn join_across_build_chunk_boundaries (
3722+ #[ values( 4 , 12 ) ] batch_size : usize ,
3723+ #[ values( 1 , 7 ) ] left_batch_rows : usize ,
3724+ #[ values(
3725+ JoinType :: Inner ,
3726+ JoinType :: Left ,
3727+ JoinType :: Right ,
3728+ JoinType :: Full ,
3729+ JoinType :: LeftSemi ,
3730+ JoinType :: LeftAnti ,
3731+ JoinType :: LeftMark ,
3732+ JoinType :: RightSemi ,
3733+ JoinType :: RightAnti ,
3734+ JoinType :: RightMark
3735+ ) ]
3736+ join_type : JoinType ,
3737+ ) -> Result < ( ) > {
3738+ let rows = CHUNKED_LEFT_ROWS as usize ;
3739+ // Batches above half the target bypass the coalescer and stay whole.
3740+ let expected_chunks = if left_batch_rows > batch_size / 2 {
3741+ rows. div_ceil ( left_batch_rows)
3742+ } else {
3743+ rows. div_ceil ( batch_size)
3744+ } ;
3745+ assert ! ( expected_chunks > 1 ) ;
3746+ let chunked_left = chunked_left_table ( Some ( left_batch_rows) ) ;
3747+ assert_eq ! (
3748+ build_side_chunk_count( & chunked_left, batch_size) . await ?,
3749+ expected_chunks
3750+ ) ;
3751+ let single_chunk_left = chunked_left_table ( None ) ;
3752+ assert_eq ! (
3753+ build_side_chunk_count( & single_chunk_left, batch_size) . await ?,
3754+ 1
3755+ ) ;
3756+
3757+ let ( columns, chunked, _) = multi_partitioned_join_collect (
3758+ chunked_left,
3759+ chunked_right_table ( ) ,
3760+ & join_type,
3761+ Some ( equality_join_filter ( ) ) ,
3762+ new_task_ctx ( batch_size) ,
3763+ )
3764+ . await ?;
3765+ let ( _, single_chunk, _) = multi_partitioned_join_collect (
3766+ single_chunk_left,
3767+ chunked_right_table ( ) ,
3768+ & join_type,
3769+ Some ( equality_join_filter ( ) ) ,
3770+ new_task_ctx ( batch_size) ,
3771+ )
3772+ . await ?;
3773+ assert_eq ! (
3774+ batches_to_sort_string( & chunked) ,
3775+ batches_to_sort_string( & single_chunk)
3776+ ) ;
3777+
3778+ if join_type == JoinType :: Left {
3779+ let a2 = columns. iter ( ) . position ( |c| c == "a2" ) . unwrap ( ) ;
3780+ let ( mut matched, mut unmatched) = ( Vec :: new ( ) , Vec :: new ( ) ) ;
3781+ for batch in & chunked {
3782+ let a1 = as_int32_array ( batch. column ( 0 ) ) ?;
3783+ for row in 0 ..batch. num_rows ( ) {
3784+ if batch. column ( a2) . is_null ( row) {
3785+ unmatched. push ( a1. value ( row) ) ;
3786+ } else {
3787+ matched. push ( a1. value ( row) ) ;
3788+ }
3789+ }
3790+ }
3791+ matched. sort_unstable ( ) ;
3792+ unmatched. sort_unstable ( ) ;
3793+ assert_eq ! ( matched, CHUNKED_LEFT_MATCHED ) ;
3794+ assert_eq ! (
3795+ unmatched,
3796+ ( 0 ..CHUNKED_LEFT_ROWS )
3797+ . filter( |i| !CHUNKED_LEFT_MATCHED . contains( i) )
3798+ . collect:: <Vec <_>>( )
3799+ ) ;
3800+ }
3801+ Ok ( ( ) )
3802+ }
3803+
36333804 /// An input that can be executed only once: later executions yield no batches, the way a
36343805 /// stream backed by an external one-shot iterator behaves.
36353806 #[ derive( Debug ) ]
0 commit comments