1use std::fmt::Formatter;
23use std::sync::Arc;
24
25use super::bitwise_stream::BitwiseSortMergeJoinStream;
26use super::materializing_stream::MaterializingSortMergeJoinStream;
27use super::metrics::SortMergeJoinMetrics;
28use crate::execution_plan::{EmissionType, boundedness_from_children};
29use crate::expressions::PhysicalSortExpr;
30use crate::joins::utils::{
31 JoinFilter, JoinOn, JoinOnRef, build_join_schema, check_join_is_valid,
32 estimate_join_statistics, reorder_output_after_swap,
33 symmetric_join_output_partitioning,
34};
35use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet, SpillMetrics};
36use crate::projection::{
37 ProjectionExec, join_allows_pushdown, join_table_borders, new_join_children,
38 physical_to_column_exprs, update_join_on,
39};
40use crate::spill::spill_manager::SpillManager;
41use crate::statistics::{ChildStats, StatisticsArgs};
42use crate::{
43 ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan,
44 ExecutionPlanProperties, InputDistributionRequirements, PlanProperties,
45 ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, validate_child_count,
46};
47
48use arrow::compute::SortOptions;
49use arrow::datatypes::SchemaRef;
50use datafusion_common::tree_node::TreeNodeRecursion;
51use datafusion_common::{
52 JoinSide, JoinType, NullEquality, Result, assert_eq_or_internal_err, internal_err,
53 plan_err,
54};
55use datafusion_execution::TaskContext;
56use datafusion_execution::memory_pool::MemoryConsumer;
57use datafusion_physical_expr::equivalence::join_equivalence_properties;
58use datafusion_physical_expr_common::physical_expr::{PhysicalExprRef, fmt_sql};
59use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements};
60
61#[derive(Debug, Clone)]
108pub struct SortMergeJoinExec {
109 pub left: Arc<dyn ExecutionPlan>,
111 pub right: Arc<dyn ExecutionPlan>,
113 pub on: JoinOn,
115 pub filter: Option<JoinFilter>,
117 pub join_type: JoinType,
119 schema: SchemaRef,
121 metrics: ExecutionPlanMetricsSet,
123 left_sort_exprs: LexOrdering,
125 right_sort_exprs: LexOrdering,
127 pub sort_options: Vec<SortOptions>,
129 pub null_equality: NullEquality,
131 cache: Arc<PlanProperties>,
133}
134
135impl SortMergeJoinExec {
136 pub fn try_new(
141 left: Arc<dyn ExecutionPlan>,
142 right: Arc<dyn ExecutionPlan>,
143 on: JoinOn,
144 filter: Option<JoinFilter>,
145 join_type: JoinType,
146 sort_options: Vec<SortOptions>,
147 null_equality: NullEquality,
148 ) -> Result<Self> {
149 let left_schema = left.schema();
150 let right_schema = right.schema();
151
152 check_join_is_valid(&left_schema, &right_schema, &on)?;
153 if sort_options.len() != on.len() {
154 return plan_err!(
155 "Expected number of sort options: {}, actual: {}",
156 on.len(),
157 sort_options.len()
158 );
159 }
160
161 let (left_sort_exprs, right_sort_exprs): (Vec<_>, Vec<_>) = on
162 .iter()
163 .zip(sort_options.iter())
164 .map(|((l, r), sort_op)| {
165 let left = PhysicalSortExpr {
166 expr: Arc::clone(l),
167 options: *sort_op,
168 };
169 let right = PhysicalSortExpr {
170 expr: Arc::clone(r),
171 options: *sort_op,
172 };
173 (left, right)
174 })
175 .unzip();
176 let Some(left_sort_exprs) = LexOrdering::new(left_sort_exprs) else {
177 return plan_err!(
178 "SortMergeJoinExec requires valid sort expressions for its left side"
179 );
180 };
181 let Some(right_sort_exprs) = LexOrdering::new(right_sort_exprs) else {
182 return plan_err!(
183 "SortMergeJoinExec requires valid sort expressions for its right side"
184 );
185 };
186
187 let schema =
188 Arc::new(build_join_schema(&left_schema, &right_schema, &join_type).0);
189 let cache =
190 Self::compute_properties(&left, &right, Arc::clone(&schema), join_type, &on)?;
191 Ok(Self {
192 left,
193 right,
194 on,
195 filter,
196 join_type,
197 schema,
198 metrics: ExecutionPlanMetricsSet::new(),
199 left_sort_exprs,
200 right_sort_exprs,
201 sort_options,
202 null_equality,
203 cache: Arc::new(cache),
204 })
205 }
206
207 pub fn probe_side(join_type: &JoinType) -> JoinSide {
210 match join_type {
213 JoinType::Right
215 | JoinType::RightSemi
216 | JoinType::RightAnti
217 | JoinType::RightMark => JoinSide::Right,
218 JoinType::Inner
219 | JoinType::Left
220 | JoinType::Full
221 | JoinType::LeftAnti
222 | JoinType::LeftSemi
223 | JoinType::LeftMark => JoinSide::Left,
224 }
225 }
226
227 fn maintains_input_order(join_type: JoinType) -> Vec<bool> {
229 match join_type {
230 JoinType::Inner => vec![true, false],
231 JoinType::Left
232 | JoinType::LeftSemi
233 | JoinType::LeftAnti
234 | JoinType::LeftMark => vec![true, false],
235 JoinType::Right
236 | JoinType::RightSemi
237 | JoinType::RightAnti
238 | JoinType::RightMark => {
239 vec![false, true]
240 }
241 _ => vec![false, false],
242 }
243 }
244
245 pub fn on(&self) -> &[(PhysicalExprRef, PhysicalExprRef)] {
247 &self.on
248 }
249
250 pub fn right(&self) -> &Arc<dyn ExecutionPlan> {
252 &self.right
253 }
254
255 pub fn join_type(&self) -> JoinType {
257 self.join_type
258 }
259
260 pub fn left(&self) -> &Arc<dyn ExecutionPlan> {
262 &self.left
263 }
264
265 pub fn filter(&self) -> &Option<JoinFilter> {
267 &self.filter
268 }
269
270 pub fn sort_options(&self) -> &[SortOptions] {
272 &self.sort_options
273 }
274
275 pub fn null_equality(&self) -> NullEquality {
277 self.null_equality
278 }
279
280 fn compute_properties(
282 left: &Arc<dyn ExecutionPlan>,
283 right: &Arc<dyn ExecutionPlan>,
284 schema: SchemaRef,
285 join_type: JoinType,
286 join_on: JoinOnRef,
287 ) -> Result<PlanProperties> {
288 let eq_properties = join_equivalence_properties(
290 left.equivalence_properties().clone(),
291 right.equivalence_properties().clone(),
292 &join_type,
293 schema,
294 &Self::maintains_input_order(join_type),
295 Some(Self::probe_side(&join_type)),
296 join_on,
297 )?;
298
299 let output_partitioning =
300 symmetric_join_output_partitioning(left, right, &join_type)?;
301
302 Ok(PlanProperties::new(
303 eq_properties,
304 output_partitioning,
305 EmissionType::Incremental,
306 boundedness_from_children([left, right]),
307 ))
308 }
309
310 pub fn swap_inputs(&self) -> Result<Arc<dyn ExecutionPlan>> {
316 let left = self.left();
317 let right = self.right();
318 let new_join = SortMergeJoinExec::try_new(
319 Arc::clone(right),
320 Arc::clone(left),
321 self.on()
322 .iter()
323 .map(|(l, r)| (Arc::clone(r), Arc::clone(l)))
324 .collect::<Vec<_>>(),
325 self.filter().as_ref().map(JoinFilter::swap),
326 self.join_type().swap(),
327 self.sort_options.clone(),
328 self.null_equality,
329 )?;
330
331 if matches!(
334 self.join_type(),
335 JoinType::LeftSemi
336 | JoinType::RightSemi
337 | JoinType::LeftAnti
338 | JoinType::RightAnti
339 | JoinType::LeftMark
340 | JoinType::RightMark
341 ) {
342 Ok(Arc::new(new_join))
343 } else {
344 reorder_output_after_swap(Arc::new(new_join), &left.schema(), &right.schema())
345 }
346 }
347}
348
349impl DisplayAs for SortMergeJoinExec {
350 fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
351 match t {
352 DisplayFormatType::Default | DisplayFormatType::Verbose => {
353 let on = self
354 .on
355 .iter()
356 .map(|(c1, c2)| format!("({c1}, {c2})"))
357 .collect::<Vec<String>>()
358 .join(", ");
359 let display_null_equality =
360 if self.null_equality() == NullEquality::NullEqualsNull {
361 ", NullsEqual: true"
362 } else {
363 ""
364 };
365 write!(
366 f,
367 "{}: join_type={:?}, on=[{}]{}{}",
368 Self::static_name(),
369 self.join_type,
370 on,
371 self.filter.as_ref().map_or_else(
372 || "".to_string(),
373 |f| format!(", filter={}", f.expression())
374 ),
375 display_null_equality,
376 )
377 }
378 DisplayFormatType::TreeRender => {
379 let on = self
380 .on
381 .iter()
382 .map(|(c1, c2)| {
383 format!("({} = {})", fmt_sql(c1.as_ref()), fmt_sql(c2.as_ref()))
384 })
385 .collect::<Vec<String>>()
386 .join(", ");
387
388 if self.join_type() != JoinType::Inner {
389 writeln!(f, "join_type={:?}", self.join_type)?;
390 }
391 writeln!(f, "on={on}")?;
392
393 if self.null_equality() == NullEquality::NullEqualsNull {
394 writeln!(f, "NullsEqual: true")?;
395 }
396
397 Ok(())
398 }
399 }
400 }
401}
402
403impl ExecutionPlan for SortMergeJoinExec {
404 fn name(&self) -> &'static str {
405 "SortMergeJoinExec"
406 }
407
408 fn properties(&self) -> &Arc<PlanProperties> {
409 &self.cache
410 }
411
412 fn required_input_distribution(&self) -> Vec<Distribution> {
413 self.input_distribution_requirements().into_per_child()
414 }
415
416 fn input_distribution_requirements(&self) -> InputDistributionRequirements {
417 let (left_expr, right_expr) = self
418 .on
419 .iter()
420 .map(|(l, r)| (Arc::clone(l), Arc::clone(r)))
421 .unzip();
422 InputDistributionRequirements::co_partitioned(vec![
423 Distribution::KeyPartitioned(left_expr),
424 Distribution::KeyPartitioned(right_expr),
425 ])
426 }
427
428 fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
429 vec![
430 Some(OrderingRequirements::from(self.left_sort_exprs.clone())),
431 Some(OrderingRequirements::from(self.right_sort_exprs.clone())),
432 ]
433 }
434
435 fn maintains_input_order(&self) -> Vec<bool> {
436 Self::maintains_input_order(self.join_type)
437 }
438
439 fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
440 vec![&self.left, &self.right]
441 }
442
443 fn apply_expressions(
444 &self,
445 f: &mut dyn FnMut(&Arc<dyn crate::PhysicalExpr>) -> Result<TreeNodeRecursion>,
446 ) -> Result<TreeNodeRecursion> {
447 let join_keys = self.on.iter().flat_map(|(left, right)| [left, right]);
448 let filter = self.filter.iter().map(|filter| filter.expression());
449 crate::apply_expression_roots(join_keys.chain(filter), f)
450 }
451
452 fn replace_children(
453 self: Arc<Self>,
454 mut children: Vec<Arc<dyn ExecutionPlan>>,
455 options: ReplaceChildrenOptions,
456 ) -> Result<Arc<dyn ExecutionPlan>> {
457 validate_child_count!(self, children);
458 match options.children_properties {
459 ChildrenPropertiesMode::Keep => {
460 let left = children.swap_remove(0);
461 let right = children.swap_remove(0);
462 Ok(Arc::new(Self {
463 left,
464 right,
465 metrics: ExecutionPlanMetricsSet::new(),
466 ..Self::clone(&*self)
467 }))
468 }
469 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 )?)),
479 _ => internal_err!("SortMergeJoin wrong number of children"),
480 },
481 }
482 }
483
484 fn with_new_children(
485 self: Arc<Self>,
486 children: Vec<Arc<dyn ExecutionPlan>>,
487 ) -> Result<Arc<dyn ExecutionPlan>> {
488 self.replace_children(
489 children,
490 ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
491 )
492 }
493
494 fn with_new_children_and_same_properties(
495 self: Arc<Self>,
496 children: Vec<Arc<dyn ExecutionPlan>>,
497 ) -> Result<Arc<dyn ExecutionPlan>> {
498 self.replace_children(
499 children,
500 ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep),
501 )
502 }
503
504 fn execute(
505 &self,
506 partition: usize,
507 context: Arc<TaskContext>,
508 ) -> Result<SendableRecordBatchStream> {
509 let left_partitions = self.left.output_partitioning().partition_count();
510 let right_partitions = self.right.output_partitioning().partition_count();
511 assert_eq_or_internal_err!(
512 left_partitions,
513 right_partitions,
514 "Invalid SortMergeJoinExec, partition count mismatch {left_partitions}!={right_partitions},\
515 consider using RepartitionExec"
516 );
517 let (on_left, on_right) = self.on.iter().cloned().unzip();
518 let (streamed, buffered, on_streamed, on_buffered) =
519 if SortMergeJoinExec::probe_side(&self.join_type) == JoinSide::Left {
520 (
521 Arc::clone(&self.left),
522 Arc::clone(&self.right),
523 on_left,
524 on_right,
525 )
526 } else {
527 (
528 Arc::clone(&self.right),
529 Arc::clone(&self.left),
530 on_right,
531 on_left,
532 )
533 };
534
535 let streamed = streamed.execute(partition, Arc::clone(&context))?;
537 let buffered = buffered.execute(partition, Arc::clone(&context))?;
538
539 let batch_size = context.session_config().batch_size();
540 let reservation = MemoryConsumer::new(format!("SMJStream[{partition}]"))
541 .register(context.memory_pool());
542 let spill_manager = SpillManager::new(
543 context.runtime_env(),
544 SpillMetrics::new(&self.metrics, partition),
545 buffered.schema(),
546 )
547 .with_compression_type(context.session_config().spill_compression());
548
549 if matches!(
550 self.join_type,
551 JoinType::LeftSemi
552 | JoinType::LeftAnti
553 | JoinType::RightSemi
554 | JoinType::RightAnti
555 | JoinType::LeftMark
556 | JoinType::RightMark
557 ) {
558 BitwiseSortMergeJoinStream::try_new(
559 Arc::clone(&self.schema),
560 self.sort_options.clone(),
561 self.null_equality,
562 streamed,
563 buffered,
564 on_streamed,
565 on_buffered,
566 self.filter.clone(),
567 self.join_type,
568 batch_size,
569 partition,
570 &self.metrics,
571 reservation,
572 spill_manager,
573 context.runtime_env(),
574 )
575 } else {
576 MaterializingSortMergeJoinStream::try_new(
577 Arc::clone(&self.schema),
578 self.sort_options.clone(),
579 self.null_equality,
580 streamed,
581 buffered,
582 on_streamed,
583 on_buffered,
584 self.filter.clone(),
585 self.join_type,
586 batch_size,
587 SortMergeJoinMetrics::new(partition, &self.metrics),
588 reservation,
589 spill_manager,
590 context.runtime_env(),
591 )
592 }
593 }
594
595 fn metrics(&self) -> Option<MetricsSet> {
596 Some(self.metrics.clone_inner())
597 }
598
599 fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {
600 vec![ChildStats::At(partition), ChildStats::At(partition)]
601 }
602
603 fn statistics_from_inputs(
604 &self,
605 input_stats: &[Arc<Statistics>],
606 _args: &StatisticsArgs,
607 ) -> Result<Arc<Statistics>> {
608 let left_stats = input_stats[0].as_ref().clone();
616 let right_stats = input_stats[1].as_ref().clone();
617 Ok(Arc::new(estimate_join_statistics(
618 left_stats,
619 right_stats,
620 &self.on,
621 self.null_equality,
622 &self.join_type,
623 &self.schema,
624 )?))
625 }
626
627 fn try_swapping_with_projection(
631 &self,
632 projection: &ProjectionExec,
633 ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
634 let Some(projection_as_columns) = physical_to_column_exprs(projection.expr())
636 else {
637 return Ok(None);
638 };
639
640 let (far_right_left_col_ind, far_left_right_col_ind) = join_table_borders(
641 self.left().schema().fields().len(),
642 &projection_as_columns,
643 );
644
645 if !join_allows_pushdown(
646 &projection_as_columns,
647 &self.schema(),
648 far_right_left_col_ind,
649 far_left_right_col_ind,
650 ) {
651 return Ok(None);
652 }
653
654 let Some(new_on) = update_join_on(
655 &projection_as_columns[0..=far_right_left_col_ind as _],
656 &projection_as_columns[far_left_right_col_ind as _..],
657 self.on(),
658 self.left().schema().fields().len(),
659 ) else {
660 return Ok(None);
661 };
662
663 let (new_left, new_right) = new_join_children(
664 &projection_as_columns,
665 far_right_left_col_ind,
666 far_left_right_col_ind,
667 self.children()[0],
668 self.children()[1],
669 )?;
670
671 Ok(Some(Arc::new(SortMergeJoinExec::try_new(
672 Arc::new(new_left),
673 Arc::new(new_right),
674 new_on,
675 self.filter.clone(),
676 self.join_type,
677 self.sort_options.clone(),
678 self.null_equality,
679 )?)))
680 }
681
682 #[cfg(feature = "proto")]
683 fn try_to_proto(
684 &self,
685 ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>,
686 ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
687 use datafusion_proto_models::protobuf;
688
689 let left = ctx.encode_child(self.left())?;
690 let right = ctx.encode_child(self.right())?;
691 let on = self
692 .on()
693 .iter()
694 .map(|(left, right)| {
695 Ok(protobuf::JoinOn {
696 left: Some(ctx.encode_expr(left)?),
697 right: Some(ctx.encode_expr(right)?),
698 })
699 })
700 .collect::<Result<Vec<_>>>()?;
701
702 let join_type = crate::joins::proto::join_type_to_proto(self.join_type());
703 let null_equality =
704 crate::joins::proto::null_equality_to_proto(self.null_equality());
705 let filter = self
706 .filter()
707 .as_ref()
708 .map(|filter| crate::joins::proto::join_filter_to_proto(filter, ctx))
709 .transpose()?;
710 let sort_options = self
711 .sort_options()
712 .iter()
713 .map(|options| protobuf::SortExprNode {
714 expr: None,
715 asc: !options.descending,
716 nulls_first: options.nulls_first,
717 })
718 .collect();
719
720 Ok(Some(protobuf::PhysicalPlanNode {
721 physical_plan_type: Some(
722 protobuf::physical_plan_node::PhysicalPlanType::SortMergeJoin(Box::new(
723 protobuf::SortMergeJoinExecNode {
724 left: Some(Box::new(left)),
725 right: Some(Box::new(right)),
726 on,
727 join_type: join_type.into(),
728 filter,
729 sort_options,
730 null_equality: null_equality.into(),
731 },
732 )),
733 ),
734 }))
735 }
736}
737
738#[cfg(feature = "proto")]
739impl SortMergeJoinExec {
740 pub fn try_from_proto(
746 node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
747 ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>,
748 ) -> Result<Arc<dyn ExecutionPlan>> {
749 use datafusion_proto_models::protobuf;
750
751 let sort_join = crate::expect_plan_variant!(
752 node,
753 protobuf::physical_plan_node::PhysicalPlanType::SortMergeJoin,
754 "SortMergeJoinExec",
755 );
756 let left = ctx.decode_required_child(
757 sort_join.left.as_deref(),
758 "SortMergeJoinExec",
759 "left",
760 )?;
761 let right = ctx.decode_required_child(
762 sort_join.right.as_deref(),
763 "SortMergeJoinExec",
764 "right",
765 )?;
766 let left_schema = left.schema();
767 let right_schema = right.schema();
768 let on = sort_join
769 .on
770 .iter()
771 .map(|columns| {
772 let left = ctx.decode_required_expr(
773 columns.left.as_ref(),
774 left_schema.as_ref(),
775 "SortMergeJoinExec",
776 "on.left",
777 )?;
778 let right = ctx.decode_required_expr(
779 columns.right.as_ref(),
780 right_schema.as_ref(),
781 "SortMergeJoinExec",
782 "on.right",
783 )?;
784 Ok((left, right))
785 })
786 .collect::<Result<JoinOn>>()?;
787
788 let join_type = crate::joins::proto::join_type_from_proto(
789 sort_join.join_type,
790 "SortMergeJoinExec",
791 )?;
792 let null_equality = crate::joins::proto::null_equality_from_proto(
793 sort_join.null_equality,
794 "SortMergeJoinExec",
795 )?;
796 let filter = sort_join
797 .filter
798 .as_ref()
799 .map(|filter| {
800 crate::joins::proto::join_filter_from_proto(
801 filter,
802 ctx,
803 "SortMergeJoinExec",
804 )
805 })
806 .transpose()?;
807 let sort_options = sort_join
808 .sort_options
809 .iter()
810 .map(|options| SortOptions {
811 descending: !options.asc,
812 nulls_first: options.nulls_first,
813 })
814 .collect();
815
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 )?))
825 }
826}