This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-25490-0b95fc29ff143397c4dccb462ce4cf69dcb33ad0 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 4ea7b6d17ef8ceea6cade0f45caa346e6735d9f5 Author: Jay Zhan <[email protected]> AuthorDate: Mon Sep 21 13:57:46 2026 +0000 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. --- .../src/joins/sort_merge_join/exec.rs | 207 +++++++++++++-------- 1 file changed, 130 insertions(+), 77 deletions(-) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index b0433250c0..c0e57a93a4 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -583,86 +583,27 @@ impl ExecutionPlan for SortMergeJoinExec { "Invalid SortMergeJoinExec, partition count mismatch {left_partitions}!={right_partitions},\ consider using RepartitionExec" ); - let (on_left, on_right) = self.on.iter().cloned().unzip(); - let (streamed, buffered, on_streamed, on_buffered) = - if SortMergeJoinExec::probe_side(&self.join_type) == JoinSide::Left { - ( - Arc::clone(&self.left), - Arc::clone(&self.right), - on_left, - on_right, - ) - } else { - ( - Arc::clone(&self.right), - Arc::clone(&self.left), - on_right, - on_left, - ) - }; - // execute children plans - let streamed = streamed.execute(partition, Arc::clone(&context))?; - let buffered = buffered.execute(partition, Arc::clone(&context))?; - - let batch_size = context.session_config().batch_size(); - // The stream spills its buffered batches when it cannot grow, so a - // pool that budgets spillable and unspillable consumers differently - // (`FairSpillPool`) has to know it can. - let reservation = MemoryConsumer::new(format!("SMJStream[{partition}]")) - .with_can_spill(true) - .register(context.memory_pool()); - let spill_manager = SpillManager::new( - context.runtime_env(), - SpillMetrics::new(&self.metrics, partition), - buffered.schema(), - ) - .with_compression_type(context.session_config().spill_compression()); + let left = self.left.execute(partition, Arc::clone(&context))?; + let right = self.right.execute(partition, Arc::clone(&context))?; + let (on_left, on_right) = self.on.iter().cloned().unzip(); - let joined = if matches!( - self.join_type, - JoinType::LeftSemi - | JoinType::LeftAnti - | JoinType::RightSemi - | JoinType::RightAnti - | JoinType::LeftMark - | JoinType::RightMark - ) { - BitwiseSortMergeJoinStream::try_new( - Arc::clone(&self.schema), - self.sort_options.clone(), - self.null_equality, - streamed, - buffered, - on_streamed, - on_buffered, - self.filter.clone(), - self.join_type, - batch_size, + let joined = sort_merge_join_stream( + SortMergeJoinInputs { + schema: Arc::clone(&self.schema), + sort_options: self.sort_options.clone(), + null_equality: self.null_equality, + left, + right, + on_left, + on_right, + filter: self.filter.clone(), + join_type: self.join_type, partition, - &self.metrics, - reservation, - spill_manager, - context.runtime_env(), - ) - } else { - MaterializingSortMergeJoinStream::try_new( - Arc::clone(&self.schema), - self.sort_options.clone(), - self.null_equality, - streamed, - buffered, - on_streamed, - on_buffered, - self.filter.clone(), - self.join_type, - batch_size, - SortMergeJoinMetrics::new(partition, &self.metrics), - reservation, - spill_manager, - context.runtime_env(), - ) - }?; + }, + &self.metrics, + &context, + )?; let Some(projection) = self.projection.clone() else { return Ok(joined); @@ -949,3 +890,115 @@ impl SortMergeJoinExec { )) } } + +/// The two sorted inputs of one partition of a sort-merge join, and how to +/// join them. +pub(crate) struct SortMergeJoinInputs { + /// The join schema (before any projection) + pub(crate) schema: SchemaRef, + /// Sort options of the join keys, one per key, that both inputs are sorted with + pub(crate) sort_options: Vec<SortOptions>, + pub(crate) null_equality: NullEquality, + /// Left input, sorted on `on_left` with `sort_options` + pub(crate) left: SendableRecordBatchStream, + /// Right input, sorted on `on_right` with `sort_options` + pub(crate) right: SendableRecordBatchStream, + pub(crate) on_left: Vec<PhysicalExprRef>, + pub(crate) on_right: Vec<PhysicalExprRef>, + pub(crate) filter: Option<JoinFilter>, + pub(crate) join_type: JoinType, + pub(crate) partition: usize, +} + +/// Joins two sorted inputs with the sort-merge join algorithm. +/// +/// Picks the streamed and buffered side by join type and the join stream +/// implementation by join type family. [`SortMergeJoinExec::execute`] is the +/// only caller today; it is a separate function so that an operator which +/// already holds two sorted streams can reuse it. The stream's metrics are +/// registered in `metrics`; its buffered side spills through the context's +/// disk manager under the memory pool's control. +pub(crate) fn sort_merge_join_stream( + inputs: SortMergeJoinInputs, + metrics: &ExecutionPlanMetricsSet, + context: &Arc<TaskContext>, +) -> Result<SendableRecordBatchStream> { + let SortMergeJoinInputs { + schema, + sort_options, + null_equality, + left, + right, + on_left, + on_right, + filter, + join_type, + partition, + } = inputs; + + let (streamed, buffered, on_streamed, on_buffered) = + if SortMergeJoinExec::probe_side(&join_type) == JoinSide::Left { + (left, right, on_left, on_right) + } else { + (right, left, on_right, on_left) + }; + + let batch_size = context.session_config().batch_size(); + // The stream spills its buffered batches when it cannot grow, so a + // pool that budgets spillable and unspillable consumers differently + // (`FairSpillPool`) has to know it can. + let reservation = MemoryConsumer::new(format!("SMJStream[{partition}]")) + .with_can_spill(true) + .register(context.memory_pool()); + let spill_manager = SpillManager::new( + context.runtime_env(), + SpillMetrics::new(metrics, partition), + buffered.schema(), + ) + .with_compression_type(context.session_config().spill_compression()); + + if matches!( + join_type, + JoinType::LeftSemi + | JoinType::LeftAnti + | JoinType::RightSemi + | JoinType::RightAnti + | JoinType::LeftMark + | JoinType::RightMark + ) { + BitwiseSortMergeJoinStream::try_new( + schema, + sort_options, + null_equality, + streamed, + buffered, + on_streamed, + on_buffered, + filter, + join_type, + batch_size, + partition, + metrics, + reservation, + spill_manager, + context.runtime_env(), + ) + } else { + MaterializingSortMergeJoinStream::try_new( + schema, + sort_options, + null_equality, + streamed, + buffered, + on_streamed, + on_buffered, + filter, + join_type, + batch_size, + SortMergeJoinMetrics::new(partition, metrics), + reservation, + spill_manager, + context.runtime_env(), + ) + } +} --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
