Skip to content

Commit cee1dc7

Browse files
authored
refactor: factor SortMergeJoinExec::execute into a reusable sort_merge_join_stream (#25490)
## Which issue does this PR close? - No separate issue. Split out of #25217 so it can be reviewed on its own. ## Rationale for this change #25217 lets a hash join that runs out of memory finish as a sort-merge join: it sorts both inputs itself and then needs to run them through the same join streams `SortMergeJoinExec` uses. Today the only way to build those streams is inside `SortMergeJoinExec::execute`, which starts from two child `ExecutionPlan`s rather than two already-sorted streams. This PR separates "join two sorted streams" from "execute my children", so the first half can be reused. It is a refactor with no new caller yet; it is split out to keep the sort-merge join changes out of the (much larger) hash join PR. ## What changes are included in this PR? All in `sort_merge_join/exec.rs`: - New `pub(crate) fn sort_merge_join_stream(SortMergeJoinInputs, &metrics, &context)`. It holds what `execute` did after executing its children: pick the streamed and buffered side by join type, register the (spillable) memory consumer, build the spill manager, and choose `BitwiseSortMergeJoinStream` or `MaterializingSortMergeJoinStream` by join type family. - New `pub(crate) struct SortMergeJoinInputs` carrying the two sorted streams, join keys, filter, join type, sort options, null equality, schema and partition. - `SortMergeJoinExec::execute` keeps the partition-count check, executes both children, calls the new function and applies the embedded projection. Nothing is `pub`, so there is no public API change. One ordering detail: `execute` used to call `execute()` on the streamed child first and the buffered child second; it now calls left then right, because the side selection moved into the shared function. For right-family joins that swaps the order of the two calls. Both still happen before anything is polled. ## What is the testing strategy for this PR? No behaviour change is intended, so this relies on the existing `sort_merge_join` unit tests, which all pass unchanged. That includes `stream_registers_as_a_spillable_consumer` and `fair_spill_pool_leaves_room_for_the_streamed_sort` from #25250, which check that the memory consumer registered by the moved code is still spillable. ## Are there any user-facing changes? No.
1 parent c607caf commit cee1dc7

1 file changed

Lines changed: 130 additions & 77 deletions

File tree

  • datafusion/physical-plan/src/joins/sort_merge_join

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

Lines changed: 130 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -583,86 +583,27 @@ impl ExecutionPlan for SortMergeJoinExec {
583583
"Invalid SortMergeJoinExec, partition count mismatch {left_partitions}!={right_partitions},\
584584
consider using RepartitionExec"
585585
);
586-
let (on_left, on_right) = self.on.iter().cloned().unzip();
587-
let (streamed, buffered, on_streamed, on_buffered) =
588-
if SortMergeJoinExec::probe_side(&self.join_type) == JoinSide::Left {
589-
(
590-
Arc::clone(&self.left),
591-
Arc::clone(&self.right),
592-
on_left,
593-
on_right,
594-
)
595-
} else {
596-
(
597-
Arc::clone(&self.right),
598-
Arc::clone(&self.left),
599-
on_right,
600-
on_left,
601-
)
602-
};
603-
604586
// execute children plans
605-
let streamed = streamed.execute(partition, Arc::clone(&context))?;
606-
let buffered = buffered.execute(partition, Arc::clone(&context))?;
607-
608-
let batch_size = context.session_config().batch_size();
609-
// The stream spills its buffered batches when it cannot grow, so a
610-
// pool that budgets spillable and unspillable consumers differently
611-
// (`FairSpillPool`) has to know it can.
612-
let reservation = MemoryConsumer::new(format!("SMJStream[{partition}]"))
613-
.with_can_spill(true)
614-
.register(context.memory_pool());
615-
let spill_manager = SpillManager::new(
616-
context.runtime_env(),
617-
SpillMetrics::new(&self.metrics, partition),
618-
buffered.schema(),
619-
)
620-
.with_compression_type(context.session_config().spill_compression());
587+
let left = self.left.execute(partition, Arc::clone(&context))?;
588+
let right = self.right.execute(partition, Arc::clone(&context))?;
589+
let (on_left, on_right) = self.on.iter().cloned().unzip();
621590

622-
let joined = if matches!(
623-
self.join_type,
624-
JoinType::LeftSemi
625-
| JoinType::LeftAnti
626-
| JoinType::RightSemi
627-
| JoinType::RightAnti
628-
| JoinType::LeftMark
629-
| JoinType::RightMark
630-
) {
631-
BitwiseSortMergeJoinStream::try_new(
632-
Arc::clone(&self.schema),
633-
self.sort_options.clone(),
634-
self.null_equality,
635-
streamed,
636-
buffered,
637-
on_streamed,
638-
on_buffered,
639-
self.filter.clone(),
640-
self.join_type,
641-
batch_size,
591+
let joined = sort_merge_join_stream(
592+
SortMergeJoinInputs {
593+
schema: Arc::clone(&self.schema),
594+
sort_options: self.sort_options.clone(),
595+
null_equality: self.null_equality,
596+
left,
597+
right,
598+
on_left,
599+
on_right,
600+
filter: self.filter.clone(),
601+
join_type: self.join_type,
642602
partition,
643-
&self.metrics,
644-
reservation,
645-
spill_manager,
646-
context.runtime_env(),
647-
)
648-
} else {
649-
MaterializingSortMergeJoinStream::try_new(
650-
Arc::clone(&self.schema),
651-
self.sort_options.clone(),
652-
self.null_equality,
653-
streamed,
654-
buffered,
655-
on_streamed,
656-
on_buffered,
657-
self.filter.clone(),
658-
self.join_type,
659-
batch_size,
660-
SortMergeJoinMetrics::new(partition, &self.metrics),
661-
reservation,
662-
spill_manager,
663-
context.runtime_env(),
664-
)
665-
}?;
603+
},
604+
&self.metrics,
605+
&context,
606+
)?;
666607

667608
let Some(projection) = self.projection.clone() else {
668609
return Ok(joined);
@@ -949,3 +890,115 @@ impl SortMergeJoinExec {
949890
))
950891
}
951892
}
893+
894+
/// The two sorted inputs of one partition of a sort-merge join, and how to
895+
/// join them.
896+
pub(crate) struct SortMergeJoinInputs {
897+
/// The join schema (before any projection)
898+
pub(crate) schema: SchemaRef,
899+
/// Sort options of the join keys, one per key, that both inputs are sorted with
900+
pub(crate) sort_options: Vec<SortOptions>,
901+
pub(crate) null_equality: NullEquality,
902+
/// Left input, sorted on `on_left` with `sort_options`
903+
pub(crate) left: SendableRecordBatchStream,
904+
/// Right input, sorted on `on_right` with `sort_options`
905+
pub(crate) right: SendableRecordBatchStream,
906+
pub(crate) on_left: Vec<PhysicalExprRef>,
907+
pub(crate) on_right: Vec<PhysicalExprRef>,
908+
pub(crate) filter: Option<JoinFilter>,
909+
pub(crate) join_type: JoinType,
910+
pub(crate) partition: usize,
911+
}
912+
913+
/// Joins two sorted inputs with the sort-merge join algorithm.
914+
///
915+
/// Picks the streamed and buffered side by join type and the join stream
916+
/// implementation by join type family. [`SortMergeJoinExec::execute`] is the
917+
/// only caller today; it is a separate function so that an operator which
918+
/// already holds two sorted streams can reuse it. The stream's metrics are
919+
/// registered in `metrics`; its buffered side spills through the context's
920+
/// disk manager under the memory pool's control.
921+
pub(crate) fn sort_merge_join_stream(
922+
inputs: SortMergeJoinInputs,
923+
metrics: &ExecutionPlanMetricsSet,
924+
context: &Arc<TaskContext>,
925+
) -> Result<SendableRecordBatchStream> {
926+
let SortMergeJoinInputs {
927+
schema,
928+
sort_options,
929+
null_equality,
930+
left,
931+
right,
932+
on_left,
933+
on_right,
934+
filter,
935+
join_type,
936+
partition,
937+
} = inputs;
938+
939+
let (streamed, buffered, on_streamed, on_buffered) =
940+
if SortMergeJoinExec::probe_side(&join_type) == JoinSide::Left {
941+
(left, right, on_left, on_right)
942+
} else {
943+
(right, left, on_right, on_left)
944+
};
945+
946+
let batch_size = context.session_config().batch_size();
947+
// The stream spills its buffered batches when it cannot grow, so a
948+
// pool that budgets spillable and unspillable consumers differently
949+
// (`FairSpillPool`) has to know it can.
950+
let reservation = MemoryConsumer::new(format!("SMJStream[{partition}]"))
951+
.with_can_spill(true)
952+
.register(context.memory_pool());
953+
let spill_manager = SpillManager::new(
954+
context.runtime_env(),
955+
SpillMetrics::new(metrics, partition),
956+
buffered.schema(),
957+
)
958+
.with_compression_type(context.session_config().spill_compression());
959+
960+
if matches!(
961+
join_type,
962+
JoinType::LeftSemi
963+
| JoinType::LeftAnti
964+
| JoinType::RightSemi
965+
| JoinType::RightAnti
966+
| JoinType::LeftMark
967+
| JoinType::RightMark
968+
) {
969+
BitwiseSortMergeJoinStream::try_new(
970+
schema,
971+
sort_options,
972+
null_equality,
973+
streamed,
974+
buffered,
975+
on_streamed,
976+
on_buffered,
977+
filter,
978+
join_type,
979+
batch_size,
980+
partition,
981+
metrics,
982+
reservation,
983+
spill_manager,
984+
context.runtime_env(),
985+
)
986+
} else {
987+
MaterializingSortMergeJoinStream::try_new(
988+
schema,
989+
sort_options,
990+
null_equality,
991+
streamed,
992+
buffered,
993+
on_streamed,
994+
on_buffered,
995+
filter,
996+
join_type,
997+
batch_size,
998+
SortMergeJoinMetrics::new(partition, metrics),
999+
reservation,
1000+
spill_manager,
1001+
context.runtime_env(),
1002+
)
1003+
}
1004+
}

0 commit comments

Comments
 (0)