Skip to content
133 changes: 132 additions & 1 deletion datafusion/core/tests/fuzz_cases/join_fuzz.rs
Original file line number Diff line number Diff line change
Expand Up @@ -662,6 +662,124 @@ async fn test_right_mark_join_1k_binary_filtered() {
}
}

// In the tests above, most key groups hold one row: the i32 tests draw the join keys
// `(a, b)` from 10,000 values for 1000 rows, and the binary ones from random bytes.
// The tests below draw them from 25 values, so each streamed row matches about 40
// buffered rows. Sort-merge join pairs a streamed row with its matches in runs, which
// end where a buffered batch ends or where the output batch fills.

#[tokio::test]
async fn test_inner_join_1k_low_cardinality() {
for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] {
JoinFuzzTestCase::new(
make_staggered_batches_i32_with_key_range(1000, 5, left_extra),
make_staggered_batches_i32_with_key_range(1000, 5, right_extra),
JoinType::Inner,
None,
)
.run_test(&[HjSmj, NljHj], false)
.await
}
}

#[tokio::test]
async fn test_inner_join_1k_low_cardinality_filtered() {
for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] {
JoinFuzzTestCase::new(
make_staggered_batches_i32_with_key_range(1000, 5, left_extra),
make_staggered_batches_i32_with_key_range(1000, 5, right_extra),
JoinType::Inner,
Some(Box::new(col_lt_col_filter)),
)
.run_test(&[HjSmj, NljHj], false)
.await
}
}

#[tokio::test]
async fn test_left_join_1k_low_cardinality() {
for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] {
JoinFuzzTestCase::new(
make_staggered_batches_i32_with_key_range(1000, 5, left_extra),
make_staggered_batches_i32_with_key_range(1000, 5, right_extra),
JoinType::Left,
None,
)
.run_test(&[HjSmj, NljHj], false)
.await
}
}

#[tokio::test]
async fn test_left_join_1k_low_cardinality_filtered() {
for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] {
JoinFuzzTestCase::new(
make_staggered_batches_i32_with_key_range(1000, 5, left_extra),
make_staggered_batches_i32_with_key_range(1000, 5, right_extra),
JoinType::Left,
Some(Box::new(col_lt_col_filter)),
)
.run_test(&[HjSmj, NljHj], false)
.await
}
}

#[tokio::test]
async fn test_right_join_1k_low_cardinality() {
for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] {
JoinFuzzTestCase::new(
make_staggered_batches_i32_with_key_range(1000, 5, left_extra),
make_staggered_batches_i32_with_key_range(1000, 5, right_extra),
JoinType::Right,
None,
)
.run_test(&[HjSmj, NljHj], false)
.await
}
}

#[tokio::test]
async fn test_right_join_1k_low_cardinality_filtered() {
for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] {
JoinFuzzTestCase::new(
make_staggered_batches_i32_with_key_range(1000, 5, left_extra),
make_staggered_batches_i32_with_key_range(1000, 5, right_extra),
JoinType::Right,
Some(Box::new(col_lt_col_filter)),
)
.run_test(&[HjSmj, NljHj], false)
.await
}
}

#[tokio::test]
async fn test_full_join_1k_low_cardinality() {
for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] {
JoinFuzzTestCase::new(
make_staggered_batches_i32_with_key_range(1000, 5, left_extra),
make_staggered_batches_i32_with_key_range(1000, 5, right_extra),
JoinType::Full,
None,
)
.run_test(&[HjSmj, NljHj], false)
.await
}
}

#[tokio::test]
async fn test_full_join_1k_low_cardinality_filtered() {
for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] {
JoinFuzzTestCase::new(
make_staggered_batches_i32_with_key_range(1000, 5, left_extra),
make_staggered_batches_i32_with_key_range(1000, 5, right_extra),
JoinType::Full,
Some(Box::new(col_lt_col_filter)),
)
.run_test(&[HjSmj, NljHj], false)
.await
}
}

type JoinFilterBuilder = Box<dyn Fn(Arc<Schema>, Arc<Schema>) -> JoinFilter>;

struct JoinFuzzTestCase {
Expand Down Expand Up @@ -1392,12 +1510,25 @@ async fn test_nested_loop_join_spill_fuzz() {
/// two sorted int32 columns 'a', 'b' ranged from 0..99 as join columns
/// two random int32 columns 'x', 'y' as other columns
fn make_staggered_batches_i32(len: usize, with_extra_column: bool) -> Vec<RecordBatch> {
make_staggered_batches_i32_with_key_range(len, 100, with_extra_column)
}

/// Like [`make_staggered_batches_i32`], but with 'a' and 'b' drawn from
/// `0..key_range`, so that a small `key_range` makes few, large key groups
fn make_staggered_batches_i32_with_key_range(
len: usize,
key_range: i32,
with_extra_column: bool,
) -> Vec<RecordBatch> {
let mut rng = rand::rng();
let mut input12: Vec<(i32, i32)> = vec![(0, 0); len];
let mut input3: Vec<i32> = vec![0; len];
let mut input4: Vec<i32> = vec![0; len];
for v in &mut input12 {
*v = (rng.random_range(0..100), rng.random_range(0..100));
*v = (
rng.random_range(0..key_range),
rng.random_range(0..key_range),
);
}
rng.fill(&mut input3[..]);
rng.fill(&mut input4[..]);
Expand Down
85 changes: 81 additions & 4 deletions datafusion/physical-plan/benches/sort_merge_join.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,11 +27,13 @@ use arrow::array::{Int64Array, RecordBatch, StringArray};
use arrow::compute::SortOptions;
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
use datafusion_common::NullEquality;
use datafusion_common::{JoinSide, NullEquality};
use datafusion_execution::TaskContext;
use datafusion_physical_expr::expressions::col;
use datafusion_expr::Operator;
use datafusion_physical_expr::expressions::{BinaryExpr, Column, col};
use datafusion_physical_plan::collect;
use datafusion_physical_plan::joins::{SortMergeJoinExec, utils::JoinOn};
use datafusion_physical_plan::joins::SortMergeJoinExec;
use datafusion_physical_plan::joins::utils::{ColumnIndex, JoinFilter, JoinOn};
use datafusion_physical_plan::test::TestMemoryExec;
use tokio::runtime::Runtime;

Expand Down Expand Up @@ -104,11 +106,45 @@ fn schema() -> SchemaRef {
]))
}

/// The join filter `left.data < right.data`.
fn data_lt_filter() -> JoinFilter {
let column_indices = vec![
ColumnIndex {
index: 1,
side: JoinSide::Left,
},
ColumnIndex {
index: 1,
side: JoinSide::Right,
},
];
let intermediate_schema = Schema::new(vec![
Field::new("left_data", DataType::Int64, false),
Field::new("right_data", DataType::Int64, false),
]);
let expression = Arc::new(BinaryExpr::new(
Arc::new(Column::new("left_data", 0)),
Operator::Lt,
Arc::new(Column::new("right_data", 1)),
));
JoinFilter::new(expression, column_indices, Arc::new(intermediate_schema))
}

fn do_join(
left: Arc<dyn datafusion_physical_plan::ExecutionPlan>,
right: Arc<dyn datafusion_physical_plan::ExecutionPlan>,
join_type: datafusion_common::JoinType,
rt: &Runtime,
) -> usize {
do_join_with_filter(left, right, join_type, None, rt)
}

fn do_join_with_filter(
left: Arc<dyn datafusion_physical_plan::ExecutionPlan>,
right: Arc<dyn datafusion_physical_plan::ExecutionPlan>,
join_type: datafusion_common::JoinType,
filter: Option<JoinFilter>,
rt: &Runtime,
) -> usize {
let on: JoinOn = vec![(
col("key", &left.schema()).unwrap(),
Expand All @@ -118,7 +154,7 @@ fn do_join(
left,
right,
on,
None,
filter,
join_type,
vec![SortOptions::default()],
NullEquality::NullEqualsNothing,
Expand Down Expand Up @@ -168,6 +204,47 @@ fn bench_smj(c: &mut Criterion) {
});
}

// 1:100 Inner Join — 20K left, 20K right, 200 distinct keys

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

inner_1to100 is the only high-fanout case, and it is an inner join. pair_streamed_row_with_group is shared by outer joins and by the deferred-filtering path, so a filtered Left or Full join with high fanout takes the same bulk path and then goes through freeze and filter. Could you add a left_1to100_filtered case, or say why the inner case is enough?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added left_1to100_filtered. It joins the inner_1to100 inputs as a left join with the filter left.data < right.data, which keeps about half of the pairs. The last streamed row of each key fails it for every buffered row, so the deferred filtering also produces null-joined rows.

Against main, as the median of seven interleaved runs, it goes from 58.57 ms to 52.86 ms (−9.8%). The bulk path saves about as much as in inner_1to100 (5.7 ms, against 6.7 ms there), so the smaller relative gain comes from the cost of the filtering, which this PR doesn't change.

// Every streamed row matches 100 buffered rows, so the join emits 2M rows.
{
let n = 20_000;
let key_mod = 200;
let left_batches = build_sorted_batches(n, key_mod, &s);
let right_batches = build_sorted_batches(n, key_mod, &s);
group.bench_function(BenchmarkId::new("inner_1to100", n), |b| {
b.iter(|| {
let left = make_exec(&left_batches, &s);
let right = make_exec(&right_batches, &s);
do_join(left, right, datafusion_common::JoinType::Inner, &rt)
})
});
}

// 1:100 Left Join with a filter — 20K left, 20K right, 200 distinct keys
// As in `inner_1to100`, every streamed row matches 100 buffered rows on the key,
// but here the pairs go through deferred filtering. `left.data < right.data`
// keeps about half of them, and none of the last streamed row of each key, which
// is null-joined.
{
let n = 20_000;
let key_mod = 200;
let left_batches = build_sorted_batches(n, key_mod, &s);
let right_batches = build_sorted_batches(n, key_mod, &s);
group.bench_function(BenchmarkId::new("left_1to100_filtered", n), |b| {
b.iter(|| {
let left = make_exec(&left_batches, &s);
let right = make_exec(&right_batches, &s);
do_join_with_filter(
left,
right,
datafusion_common::JoinType::Left,
Some(data_lt_filter()),
&rt,
)
})
});
}

// Left Join — 100K each, ~5% unmatched on left
{
let n = 100_000;
Expand Down
Loading
Loading