@@ -34,29 +34,35 @@ use crate::joins::utils::{
3434} ;
3535use crate :: metrics:: { ExecutionPlanMetricsSet , MetricsSet , SpillMetrics } ;
3636use 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} ;
4040use crate :: spill:: spill_manager:: SpillManager ;
4141use crate :: statistics:: { ChildStats , StatisticsArgs } ;
42+ use crate :: stream:: RecordBatchStreamAdapter ;
4243use 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
4850use arrow:: compute:: SortOptions ;
4951use arrow:: datatypes:: SchemaRef ;
52+ use datafusion_common:: project_schema;
5053use datafusion_common:: tree_node:: TreeNodeRecursion ;
5154use datafusion_common:: {
5255 JoinSide , JoinType , NullEquality , Result , assert_eq_or_internal_err, internal_err,
5356 plan_err,
5457} ;
5558use datafusion_execution:: TaskContext ;
5659use 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+ } ;
5863use datafusion_physical_expr_common:: physical_expr:: { PhysicalExprRef , fmt_sql} ;
5964use 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+
349397impl 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