Skip to content

Commit 65415ac

Browse files
Dandandanclaude
andcommitted
feat: add projection support to SortMergeJoinExec
A sort merge join emitted every column it joined, so a query selecting a subset needed a ProjectionExec above it, and one joining on an expression kept that expression in the output. Projection pushdown could only push into the children when each side's columns stayed together, and gave up otherwise. SortMergeJoinExec now carries an optional projection, as HashJoinExec does, and projection pushdown embeds into the join whatever it cannot push into the children. Co-Authored-By: Claude Opus 5 <[email protected]> Claude-Session: https://claude.ai/code/session_012YiAABcW4WSqij31zz2P6c
1 parent 5091b42 commit 65415ac

9 files changed

Lines changed: 273 additions & 61 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+
}

‎datafusion/physical-plan/src/joins/sort_merge_join/exec.rs‎

Lines changed: 139 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -34,29 +34,35 @@ use crate::joins::utils::{
3434
};
3535
use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet, SpillMetrics};
3636
use crate::projection::{
37-
ProjectionExec, join_allows_pushdown, join_table_borders, new_join_children,
38-
physical_to_column_exprs, update_join_on,
37+
EmbeddedProjection, ProjectionExec, join_allows_pushdown, join_table_borders,
38+
new_join_children, physical_to_column_exprs, try_embed_projection, update_join_on,
3939
};
4040
use crate::spill::spill_manager::SpillManager;
4141
use crate::statistics::{ChildStats, StatisticsArgs};
42+
use crate::stream::RecordBatchStreamAdapter;
4243
use crate::{
4344
ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan,
4445
ExecutionPlanProperties, InputDistributionRequirements, PlanProperties,
45-
ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, validate_child_count,
46+
ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, common::can_project,
47+
validate_child_count,
4648
};
4749

4850
use arrow::compute::SortOptions;
4951
use arrow::datatypes::SchemaRef;
52+
use datafusion_common::project_schema;
5053
use datafusion_common::tree_node::TreeNodeRecursion;
5154
use datafusion_common::{
5255
JoinSide, JoinType, NullEquality, Result, assert_eq_or_internal_err, internal_err,
5356
plan_err,
5457
};
5558
use datafusion_execution::TaskContext;
5659
use datafusion_execution::memory_pool::MemoryConsumer;
57-
use datafusion_physical_expr::equivalence::join_equivalence_properties;
60+
use datafusion_physical_expr::equivalence::{
61+
ProjectionMapping, join_equivalence_properties,
62+
};
5863
use datafusion_physical_expr_common::physical_expr::{PhysicalExprRef, fmt_sql};
5964
use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements};
65+
use futures::StreamExt;
6066

6167
/// Join execution plan that executes equi-join predicates on multiple partitions using Sort-Merge
6268
/// join algorithm and applies an optional filter post join. Can be used to join arbitrarily large
@@ -128,6 +134,8 @@ pub struct SortMergeJoinExec {
128134
pub sort_options: Vec<SortOptions>,
129135
/// Defines the null equality for the join.
130136
pub null_equality: NullEquality,
137+
/// The columns of `schema` to emit, in order. `None` emits all of them.
138+
pub projection: Option<Vec<usize>>,
131139
/// Cache holding plan properties like equivalences, output partitioning etc.
132140
cache: Arc<PlanProperties>,
133141
}
@@ -187,7 +195,7 @@ impl SortMergeJoinExec {
187195
let schema =
188196
Arc::new(build_join_schema(&left_schema, &right_schema, &join_type).0);
189197
let cache =
190-
Self::compute_properties(&left, &right, Arc::clone(&schema), join_type, &on)?;
198+
Self::compute_properties(&left, &right, &schema, join_type, &on, None)?;
191199
Ok(Self {
192200
left,
193201
right,
@@ -200,10 +208,36 @@ impl SortMergeJoinExec {
200208
right_sort_exprs,
201209
sort_options,
202210
null_equality,
211+
projection: None,
212+
cache: Arc::new(cache),
213+
})
214+
}
215+
216+
/// Returns this join emitting only the columns in `projection`, in that order.
217+
/// The indices address the join's own schema, before any projection.
218+
pub fn with_projection(&self, projection: Option<Vec<usize>>) -> Result<Self> {
219+
can_project(&self.schema, projection.as_deref())?;
220+
let cache = Self::compute_properties(
221+
&self.left,
222+
&self.right,
223+
&self.schema,
224+
self.join_type,
225+
&self.on,
226+
projection.as_deref(),
227+
)?;
228+
Ok(Self {
229+
projection,
230+
metrics: ExecutionPlanMetricsSet::new(),
203231
cache: Arc::new(cache),
232+
..Self::clone(self)
204233
})
205234
}
206235

236+
/// Whether the join emits fewer or reordered columns than it joins.
237+
pub fn contains_projection(&self) -> bool {
238+
self.projection.is_some()
239+
}
240+
207241
/// Get probe side (e.g streaming side) information for this sort merge join.
208242
/// In current implementation, probe side is determined according to join type.
209243
pub fn probe_side(join_type: &JoinType) -> JoinSide {
@@ -281,24 +315,32 @@ impl SortMergeJoinExec {
281315
fn compute_properties(
282316
left: &Arc<dyn ExecutionPlan>,
283317
right: &Arc<dyn ExecutionPlan>,
284-
schema: SchemaRef,
318+
schema: &SchemaRef,
285319
join_type: JoinType,
286320
join_on: JoinOnRef,
321+
projection: Option<&[usize]>,
287322
) -> Result<PlanProperties> {
288323
// Calculate equivalence properties:
289-
let eq_properties = join_equivalence_properties(
324+
let mut eq_properties = join_equivalence_properties(
290325
left.equivalence_properties().clone(),
291326
right.equivalence_properties().clone(),
292327
&join_type,
293-
schema,
328+
Arc::clone(schema),
294329
&Self::maintains_input_order(join_type),
295330
Some(Self::probe_side(&join_type)),
296331
join_on,
297332
)?;
298333

299-
let output_partitioning =
334+
let mut output_partitioning =
300335
symmetric_join_output_partitioning(left, right, &join_type)?;
301336

337+
if let Some(projection) = projection {
338+
let mapping = ProjectionMapping::from_indices(projection, schema)?;
339+
let projected = project_schema(schema, Some(projection))?;
340+
output_partitioning = output_partitioning.project(&mapping, &eq_properties);
341+
eq_properties = eq_properties.project(&mapping, projected);
342+
}
343+
302344
Ok(PlanProperties::new(
303345
eq_properties,
304346
output_partitioning,
@@ -346,6 +388,12 @@ impl SortMergeJoinExec {
346388
}
347389
}
348390

391+
impl EmbeddedProjection for SortMergeJoinExec {
392+
fn with_projection(&self, projection: Option<Vec<usize>>) -> Result<Self> {
393+
self.with_projection(projection)
394+
}
395+
}
396+
349397
impl DisplayAs for SortMergeJoinExec {
350398
fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
351399
match t {
@@ -362,9 +410,27 @@ impl DisplayAs for SortMergeJoinExec {
362410
} else {
363411
""
364412
};
413+
let display_projections = if self.contains_projection() {
414+
format!(
415+
", projection=[{}]",
416+
self.projection
417+
.as_ref()
418+
.unwrap()
419+
.iter()
420+
.map(|index| format!(
421+
"{}@{}",
422+
self.schema.fields().get(*index).unwrap().name(),
423+
index
424+
))
425+
.collect::<Vec<_>>()
426+
.join(", ")
427+
)
428+
} else {
429+
"".to_string()
430+
};
365431
write!(
366432
f,
367-
"{}: join_type={:?}, on=[{}]{}{}",
433+
"{}: join_type={:?}, on=[{}]{}{}{}",
368434
Self::static_name(),
369435
self.join_type,
370436
on,
@@ -373,6 +439,7 @@ impl DisplayAs for SortMergeJoinExec {
373439
|f| format!(", filter={}", f.expression())
374440
),
375441
display_null_equality,
442+
display_projections,
376443
)
377444
}
378445
DisplayFormatType::TreeRender => {
@@ -467,15 +534,18 @@ impl ExecutionPlan for SortMergeJoinExec {
467534
}))
468535
}
469536
ChildrenPropertiesMode::Recompute => match &children[..] {
470-
[left, right] => Ok(Arc::new(SortMergeJoinExec::try_new(
471-
Arc::clone(left),
472-
Arc::clone(right),
473-
self.on.clone(),
474-
self.filter.clone(),
475-
self.join_type,
476-
self.sort_options.clone(),
477-
self.null_equality,
478-
)?)),
537+
[left, right] => Ok(Arc::new(
538+
SortMergeJoinExec::try_new(
539+
Arc::clone(left),
540+
Arc::clone(right),
541+
self.on.clone(),
542+
self.filter.clone(),
543+
self.join_type,
544+
self.sort_options.clone(),
545+
self.null_equality,
546+
)?
547+
.with_projection(self.projection.clone())?,
548+
)),
479549
_ => internal_err!("SortMergeJoin wrong number of children"),
480550
},
481551
}
@@ -546,7 +616,7 @@ impl ExecutionPlan for SortMergeJoinExec {
546616
)
547617
.with_compression_type(context.session_config().spill_compression());
548618

549-
if matches!(
619+
let joined = if matches!(
550620
self.join_type,
551621
JoinType::LeftSemi
552622
| JoinType::LeftAnti
@@ -589,7 +659,16 @@ impl ExecutionPlan for SortMergeJoinExec {
589659
spill_manager,
590660
context.runtime_env(),
591661
)
592-
}
662+
}?;
663+
664+
let Some(projection) = self.projection.clone() else {
665+
return Ok(joined);
666+
};
667+
let schema = self.schema();
668+
Ok(Box::pin(RecordBatchStreamAdapter::new(
669+
Arc::clone(&schema),
670+
joined.map(move |batch| Ok(batch?.project(&projection)?)),
671+
)))
593672
}
594673

595674
fn metrics(&self) -> Option<MetricsSet> {
@@ -614,14 +693,18 @@ impl ExecutionPlan for SortMergeJoinExec {
614693
// - `A LEFT JOIN B ON A.col=B.col` with `COUNT_DISTINCT(B.col)=COUNT(B.col)`
615694
let left_stats = input_stats[0].as_ref().clone();
616695
let right_stats = input_stats[1].as_ref().clone();
617-
Ok(Arc::new(estimate_join_statistics(
696+
let stats = estimate_join_statistics(
618697
left_stats,
619698
right_stats,
620699
&self.on,
621700
self.null_equality,
622701
&self.join_type,
623702
&self.schema,
624-
)?))
703+
)?;
704+
Ok(Arc::new(match &self.projection {
705+
Some(projection) => stats.project(Some(projection)),
706+
None => stats,
707+
}))
625708
}
626709

627710
/// Tries to swap the projection with its input [`SortMergeJoinExec`]. If it can be done,
@@ -631,6 +714,9 @@ impl ExecutionPlan for SortMergeJoinExec {
631714
&self,
632715
projection: &ProjectionExec,
633716
) -> Result<Option<Arc<dyn ExecutionPlan>>> {
717+
if self.contains_projection() {
718+
return Ok(None);
719+
}
634720
// Convert projected PhysicalExpr's to columns. If not possible, we cannot proceed.
635721
let Some(projection_as_columns) = physical_to_column_exprs(projection.expr())
636722
else {
@@ -642,13 +728,15 @@ impl ExecutionPlan for SortMergeJoinExec {
642728
&projection_as_columns,
643729
);
644730

731+
// Pushing into the children needs each side's columns to stay together, which
732+
// an arbitrary projection does not. The join can apply that one itself.
645733
if !join_allows_pushdown(
646734
&projection_as_columns,
647735
&self.schema(),
648736
far_right_left_col_ind,
649737
far_left_right_col_ind,
650738
) {
651-
return Ok(None);
739+
return try_embed_projection(projection, self);
652740
}
653741

654742
let Some(new_on) = update_join_on(
@@ -657,7 +745,7 @@ impl ExecutionPlan for SortMergeJoinExec {
657745
self.on(),
658746
self.left().schema().fields().len(),
659747
) else {
660-
return Ok(None);
748+
return try_embed_projection(projection, self);
661749
};
662750

663751
let (new_left, new_right) = new_join_children(
@@ -728,6 +816,12 @@ impl ExecutionPlan for SortMergeJoinExec {
728816
filter,
729817
sort_options,
730818
null_equality: null_equality.into(),
819+
projection: self
820+
.projection
821+
.iter()
822+
.flatten()
823+
.map(|index| *index as u32)
824+
.collect(),
731825
},
732826
)),
733827
),
@@ -813,14 +907,25 @@ impl SortMergeJoinExec {
813907
})
814908
.collect();
815909

816-
Ok(Arc::new(Self::try_new(
817-
left,
818-
right,
819-
on,
820-
filter,
821-
join_type,
822-
sort_options,
823-
null_equality,
824-
)?))
910+
let projection = (!sort_join.projection.is_empty()).then(|| {
911+
sort_join
912+
.projection
913+
.iter()
914+
.map(|index| *index as usize)
915+
.collect()
916+
});
917+
918+
Ok(Arc::new(
919+
Self::try_new(
920+
left,
921+
right,
922+
on,
923+
filter,
924+
join_type,
925+
sort_options,
926+
null_equality,
927+
)?
928+
.with_projection(projection)?,
929+
))
825930
}
826931
}

0 commit comments

Comments
 (0)