Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
207 changes: 204 additions & 3 deletions datafusion/core/src/physical_planner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1473,15 +1473,29 @@ impl DefaultPhysicalPlanner {
&& session_state.config().repartition_joins()
&& !prefer_hash_join
{
// Use SortMergeJoin if hash join is not preferred
let join_on_len = join_on.len();
// Derive sort options from the left input's existing ordering
// rather than hardcoding SortOptions::default()
let sort_options: Vec<SortOptions> = join_on
.iter()
.map(|(left_col, _)| {
physical_left
.output_ordering()
.and_then(|ordering| {
ordering
.iter()
.find(|sort_expr| sort_expr.expr.eq(left_col))
.map(|sort_expr| sort_expr.options)
})
.unwrap_or_default()
})
.collect();
Arc::new(SortMergeJoinExec::try_new(
physical_left,
physical_right,
join_on,
join_filter,
*join_type,
vec![SortOptions::default(); join_on_len],
sort_options,
*null_equality,
)?)
} else if session_state.config().target_partitions() > 1
Expand Down Expand Up @@ -4356,4 +4370,191 @@ digraph {
assert_contains!(&err_str, "field nullability at index");
assert_contains!(&err_str, "field metadata at index");
}

/// Helper: make a session state configured to use SortMergeJoin
fn make_smj_session_state() -> SessionState {
let runtime = Arc::new(RuntimeEnv::default());
let config = SessionConfig::new()
.with_target_partitions(4)
.set_bool("datafusion.optimizer.prefer_hash_join", false)
.set_bool("datafusion.optimizer.skip_failed_rules", false);
SessionStateBuilder::new()
.with_config(config)
.with_runtime_env(runtime)
.with_default_features()
.build()
}

/// Helper: create a fully-optimized physical plan and return its
/// displayable string for snapshot testing
async fn plan_smj_display(ctx: &SessionContext, sql: &str) -> Result<String> {
let state = ctx.state();
let logical_plan = state.create_logical_plan(sql).await?;
let logical_plan = state.optimize(&logical_plan)?;
let planner = DefaultPhysicalPlanner::default();
let plan = planner.create_physical_plan(&logical_plan, &state).await?;
Ok(format!("{}", displayable(plan.as_ref()).indent(false)))
}

/// Test: left child sorted DESC on join key -> SMJ derives DESC sort
/// options, avoiding a re-sort on the left side
#[tokio::test]
async fn test_smj_sort_options_derived_from_left_desc() -> Result<()> {
let schema = Arc::new(Schema::new(vec![
Field::new("a", DataType::Int32, false),
Field::new("b", DataType::Int32, false),
]));
let batch = RecordBatch::try_new(
Arc::clone(&schema),
vec![
Arc::new(Int32Array::from(vec![3, 2, 1])),
Arc::new(Int32Array::from(vec![10, 20, 30])),
],
)?;

let left = MemTable::try_new(Arc::clone(&schema), vec![vec![batch.clone()]])?
.with_sort_order(vec![vec![
col("a").sort(false, false), // DESC, nulls last
]]);
let right = MemTable::try_new(Arc::clone(&schema), vec![vec![batch]])?;

let state = make_smj_session_state();
let ctx = SessionContext::new_with_state(state);
ctx.register_table("left_table", Arc::new(left))?;
ctx.register_table("right_table", Arc::new(right))?;

let plan = plan_smj_display(
&ctx,
"SELECT * FROM left_table JOIN right_table ON left_table.a = right_table.a",
)
.await?;
insta::assert_snapshot!(plan, @r"
SortMergeJoinExec: join_type=Inner, on=[(a@0, a@0)]
DataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0 DESC NULLS LAST
SortExec: expr=[a@0 DESC NULLS LAST], preserve_partitioning=[false]
DataSourceExec: partitions=1, partition_sizes=[1]
");
Ok(())
}

/// Test: left child not sorted on join key -> SMJ falls back to default
/// ASC behavior
#[tokio::test]
async fn test_smj_sort_options_default_when_left_unsorted() -> Result<()> {
let schema = Arc::new(Schema::new(vec![
Field::new("a", DataType::Int32, false),
Field::new("b", DataType::Int32, false),
]));
let batch = RecordBatch::try_new(
Arc::clone(&schema),
vec![
Arc::new(Int32Array::from(vec![1, 2, 3])),
Arc::new(Int32Array::from(vec![10, 20, 30])),
],
)?;

let left = MemTable::try_new(Arc::clone(&schema), vec![vec![batch.clone()]])?;
let right = MemTable::try_new(Arc::clone(&schema), vec![vec![batch]])?;

let state = make_smj_session_state();
let ctx = SessionContext::new_with_state(state);
ctx.register_table("left_table", Arc::new(left))?;
ctx.register_table("right_table", Arc::new(right))?;

let plan = plan_smj_display(
&ctx,
"SELECT * FROM left_table JOIN right_table ON left_table.a = right_table.a",
)
.await?;
insta::assert_snapshot!(plan, @r"
SortMergeJoinExec: join_type=Inner, on=[(a@0, a@0)]
SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
DataSourceExec: partitions=1, partition_sizes=[1]
SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
DataSourceExec: partitions=1, partition_sizes=[1]
");
Ok(())
}

/// Test: multi-column join where left child's ordering partially matches
#[tokio::test]
async fn test_smj_sort_options_partial_match_multi_column() -> Result<()> {
let schema = Arc::new(Schema::new(vec![
Field::new("a", DataType::Int32, false),
Field::new("b", DataType::Int32, false),
]));
let batch = RecordBatch::try_new(
Arc::clone(&schema),
vec![
Arc::new(Int32Array::from(vec![1, 2, 3])),
Arc::new(Int32Array::from(vec![30, 20, 10])),
],
)?;

// Left sorted by "a" DESC only; "b" is not in the ordering
let left = MemTable::try_new(Arc::clone(&schema), vec![vec![batch.clone()]])?
.with_sort_order(vec![vec![
col("a").sort(false, false), // DESC, nulls last
]]);
let right = MemTable::try_new(Arc::clone(&schema), vec![vec![batch]])?;

let state = make_smj_session_state();
let ctx = SessionContext::new_with_state(state);
ctx.register_table("left_table", Arc::new(left))?;
ctx.register_table("right_table", Arc::new(right))?;

let plan = plan_smj_display(
&ctx,
"SELECT * FROM left_table JOIN right_table \
ON left_table.a = right_table.a AND left_table.b = right_table.b",
)
.await?;
insta::assert_snapshot!(plan, @r"
SortMergeJoinExec: join_type=Inner, on=[(a@0, a@0), (b@1, b@1)]
SortExec: expr=[a@0 DESC NULLS LAST, b@1 ASC], preserve_partitioning=[false]
DataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0 DESC NULLS LAST
SortExec: expr=[a@0 DESC NULLS LAST, b@1 ASC], preserve_partitioning=[false]
DataSourceExec: partitions=1, partition_sizes=[1]
");
Ok(())
}

/// Test: both sides already sorted DESC on join key -> no SortExec needed
#[tokio::test]
async fn test_smj_sort_options_both_sides_sorted_desc() -> Result<()> {
let schema = Arc::new(Schema::new(vec![
Field::new("a", DataType::Int32, false),
Field::new("b", DataType::Int32, false),
]));
let batch = RecordBatch::try_new(
Arc::clone(&schema),
vec![
Arc::new(Int32Array::from(vec![3, 2, 1])),
Arc::new(Int32Array::from(vec![10, 20, 30])),
],
)?;

let sort_order = vec![vec![col("a").sort(false, false)]]; // DESC, nulls last
let left = MemTable::try_new(Arc::clone(&schema), vec![vec![batch.clone()]])?
.with_sort_order(sort_order.clone());
let right = MemTable::try_new(Arc::clone(&schema), vec![vec![batch]])?
.with_sort_order(sort_order);

let state = make_smj_session_state();
let ctx = SessionContext::new_with_state(state);
ctx.register_table("left_table", Arc::new(left))?;
ctx.register_table("right_table", Arc::new(right))?;

let plan = plan_smj_display(
&ctx,
"SELECT * FROM left_table JOIN right_table ON left_table.a = right_table.a",
)
.await?;
insta::assert_snapshot!(plan, @r"
SortMergeJoinExec: join_type=Inner, on=[(a@0, a@0)]
DataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0 DESC NULLS LAST
DataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0 DESC NULLS LAST
");
Ok(())
}
}
Loading