Skip to content

Commit a5bbd22

Browse files
committed
test the NestedLoopJoin probe and unmatched-left emission across build chunk boundaries
1 parent 0e2df27 commit a5bbd22

1 file changed

Lines changed: 171 additions & 0 deletions

File tree

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

Lines changed: 171 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)