Skip to content

Commit 5134a1a

Browse files
Dandandanclaude
andauthored
feat: add projection support to SortMergeJoinExec (#24517)
## Which issue does this PR close? - Closes #24518 ## Rationale for this change `SortMergeJoinExec` emits every column of both inputs. We can add projection support to only produce the output columns (saving some compute / copy). ```sql select * from t1 right join t2 on t1.c3 = t2.c3 -- c3 needs a cast ``` ```text ProjectionExec: expr=[c1@0 as c1, ..., c4@8 as c4] SortMergeJoinExec: join_type=Right, on=[(CAST(t1.c3 AS Decimal128(10, 2))@4, c3@2)] ``` The cast column is only there to be joined on, and every operator above the join carries it until the projection removes it. ## What changes are included in this PR? `SortMergeJoinExec` takes an optional projection, set with `with_projection`, the same shape as `HashJoinExec`'s. `try_swapping_with_projection` embeds the projection into the join when it cannot push it into the children, so the query above becomes: ```text SortMergeJoinExec: join_type=Right, on=[(CAST(t1.c3 AS Decimal128(10, 2))@4, c3@2)], projection=[c1@0, c2@1, c3@2, c4@3, c1@5, c2@6, c3@7, c4@8] ``` The change doesn't bring a big speedup but helps aligning with other join types and helping simplify join optimization in other areas (e.g. join enumeration). ## Are these changes tested? Yes: - a projection pushdown test for a projection that interleaves the two sides, which the existing pushdown cannot handle - a serialization round trip - existing sqllogictests, whose plans lose a `ProjectionExec` in four places ## Are there any user-facing changes? Explain / proto changes. --------- Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
1 parent 3931ab4 commit 5134a1a

10 files changed

Lines changed: 373 additions & 67 deletions

File tree

‎datafusion/core/tests/physical_optimizer/projection_pushdown.rs‎

Lines changed: 45 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,8 +49,8 @@ use datafusion_physical_plan::coop::CooperativeExec;
4949
use datafusion_physical_plan::filter::{FilterExec, FilterExecBuilder};
5050
use datafusion_physical_plan::joins::utils::{ColumnIndex, JoinFilter};
5151
use datafusion_physical_plan::joins::{
52-
HashJoinExec, NestedLoopJoinExec, PartitionMode, StreamJoinPartitionMode,
53-
SymmetricHashJoinExec,
52+
HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec,
53+
StreamJoinPartitionMode, SymmetricHashJoinExec,
5454
};
5555
use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr, update_expr};
5656
use datafusion_physical_plan::repartition::RepartitionExec;
@@ -1875,3 +1875,46 @@ fn test_filter_with_embedded_projection_after_renaming_projection() -> Result<()
18751875

18761876
Ok(())
18771877
}
1878+
1879+
#[test]
1880+
fn test_sort_merge_join_interleaved_projection_embeds() -> Result<()> {
1881+
// SELECT t1.c, t2.c, t1.a FROM t1 JOIN t2 ON t1.b = t2.c
1882+
// Taking a column from each side in turn leaves the sides interleaved, which
1883+
// cannot be pushed into the children, so the join applies it itself.
1884+
let join = Arc::new(SortMergeJoinExec::try_new(
1885+
create_simple_csv_exec(),
1886+
create_simple_csv_exec(),
1887+
vec![(Arc::new(Column::new("b", 1)), Arc::new(Column::new("c", 2)))],
1888+
None,
1889+
JoinType::Inner,
1890+
vec![SortOptions::default()],
1891+
NullEquality::NullEqualsNothing,
1892+
)?);
1893+
let projection: Arc<dyn ExecutionPlan> = Arc::new(ProjectionExec::try_new(
1894+
vec![
1895+
ProjectionExpr::new(Arc::new(Column::new("c", 2)), "c_from_left"),
1896+
ProjectionExpr::new(Arc::new(Column::new("c", 7)), "c_from_right"),
1897+
ProjectionExpr::new(Arc::new(Column::new("a", 0)), "a_from_left"),
1898+
],
1899+
join,
1900+
)?);
1901+
1902+
let after_optimize =
1903+
ProjectionPushdown::new().optimize(projection, &ConfigOptions::new())?;
1904+
let actual = displayable(after_optimize.as_ref())
1905+
.indent(true)
1906+
.to_string()
1907+
.trim()
1908+
.to_string();
1909+
assert_snapshot!(
1910+
actual,
1911+
@r"
1912+
ProjectionExec: expr=[c@0 as c_from_left, c@1 as c_from_right, a@2 as a_from_left]
1913+
SortMergeJoinExec: join_type=Inner, on=[(b@1, c@2)], projection=[c@2, c@7, a@0]
1914+
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false
1915+
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false
1916+
"
1917+
);
1918+
1919+
Ok(())
1920+
}

0 commit comments

Comments
 (0)