This is an automated email from the ASF dual-hosted git repository.
dheres pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 7e0738a979 Remove CoalescePartitions insertion from HashJoinExec
(#15476)
7e0738a979 is described below
commit 7e0738a979054e95d1935d6e3e58b27385679031
Author: Christian <[email protected]>
AuthorDate: Sat Mar 29 18:35:02 2025 +0100
Remove CoalescePartitions insertion from HashJoinExec (#15476)
* refactor: Catch PartitionMode:Auto in execute explicitly
* refactor(hash_join): Move coalesce logic into function
* refactor(hash_join): Move coalesce logic out of collect_left_input
* refactor(hash_join): Execute build side earlier
* chore(hash_join): Drop unnecessary clippy hint \o/
* Back out "refactor: Catch PartitionMode:Auto in execute explicitly"
This backs out commit 4346cb868c95bfde3fb23351836157524a06050b.
* Remove CoalescePartitions from CollectLeft-HashJoins
* Fix imports
* Fix tests
* Check CollectLeft-HJ distribution in execute
---
datafusion/physical-plan/src/joins/hash_join.rs | 70 ++++++++++++-------------
1 file changed, 34 insertions(+), 36 deletions(-)
diff --git a/datafusion/physical-plan/src/joins/hash_join.rs
b/datafusion/physical-plan/src/joins/hash_join.rs
index c2a313edd1..d6c0a682c2 100644
--- a/datafusion/physical-plan/src/joins/hash_join.rs
+++ b/datafusion/physical-plan/src/joins/hash_join.rs
@@ -40,7 +40,6 @@ use crate::projection::{
use crate::spill::get_record_batch_memory_size;
use crate::ExecutionPlanProperties;
use crate::{
- coalesce_partitions::CoalescePartitionsExec,
common::can_project,
handle_state,
hash_utils::create_hashes,
@@ -792,34 +791,44 @@ impl ExecutionPlan for HashJoinExec {
);
}
+ if self.mode == PartitionMode::CollectLeft && left_partitions != 1 {
+ return internal_err!(
+ "Invalid HashJoinExec,the output partition count of the left
child must be 1 in CollectLeft mode,\
+ consider using CoalescePartitionsExec"
+ );
+ }
+
let join_metrics = BuildProbeJoinMetrics::new(partition,
&self.metrics);
let left_fut = match self.mode {
- PartitionMode::CollectLeft => self.left_fut.once(|| {
- let reservation =
-
MemoryConsumer::new("HashJoinInput").register(context.memory_pool());
- collect_left_input(
- None,
- self.random_state.clone(),
- Arc::clone(&self.left),
- on_left.clone(),
- Arc::clone(&context),
- join_metrics.clone(),
- reservation,
- need_produce_result_in_final(self.join_type),
- self.right().output_partitioning().partition_count(),
- )
- }),
+ PartitionMode::CollectLeft => {
+ let left_stream = self.left.execute(0, Arc::clone(&context))?;
+
+ self.left_fut.once(|| {
+ let reservation = MemoryConsumer::new("HashJoinInput")
+ .register(context.memory_pool());
+
+ collect_left_input(
+ self.random_state.clone(),
+ left_stream,
+ on_left.clone(),
+ join_metrics.clone(),
+ reservation,
+ need_produce_result_in_final(self.join_type),
+ self.right().output_partitioning().partition_count(),
+ )
+ })
+ }
PartitionMode::Partitioned => {
+ let left_stream = self.left.execute(partition,
Arc::clone(&context))?;
+
let reservation =
MemoryConsumer::new(format!("HashJoinInput[{partition}]"))
.register(context.memory_pool());
OnceFut::new(collect_left_input(
- Some(partition),
self.random_state.clone(),
- Arc::clone(&self.left),
+ left_stream,
on_left.clone(),
- Arc::clone(&context),
join_metrics.clone(),
reservation,
need_produce_result_in_final(self.join_type),
@@ -930,36 +939,22 @@ impl ExecutionPlan for HashJoinExec {
/// Reads the left (build) side of the input, buffering it in memory, to build
a
/// hash table (`LeftJoinData`)
-#[allow(clippy::too_many_arguments)]
async fn collect_left_input(
- partition: Option<usize>,
random_state: RandomState,
- left: Arc<dyn ExecutionPlan>,
+ left_stream: SendableRecordBatchStream,
on_left: Vec<PhysicalExprRef>,
- context: Arc<TaskContext>,
metrics: BuildProbeJoinMetrics,
reservation: MemoryReservation,
with_visited_indices_bitmap: bool,
probe_threads_count: usize,
) -> Result<JoinLeftData> {
- let schema = left.schema();
-
- let (left_input, left_input_partition) = if let Some(partition) =
partition {
- (left, partition)
- } else if left.output_partitioning().partition_count() != 1 {
- (Arc::new(CoalescePartitionsExec::new(left)) as _, 0)
- } else {
- (left, 0)
- };
-
- // Depending on partition argument load single partition or whole left
side in memory
- let stream = left_input.execute(left_input_partition,
Arc::clone(&context))?;
+ let schema = left_stream.schema();
// This operation performs 2 steps at once:
// 1. creates a [JoinHashMap] of all batches from the stream
// 2. stores the batches in a vector.
let initial = (Vec::new(), 0, metrics, reservation);
- let (batches, num_rows, metrics, mut reservation) = stream
+ let (batches, num_rows, metrics, mut reservation) = left_stream
.try_fold(initial, |mut acc, batch| async {
let batch_size = get_record_batch_memory_size(&batch);
// Reserve memory for incoming batch
@@ -1655,6 +1650,7 @@ impl EmbeddedProjection for HashJoinExec {
#[cfg(test)]
mod tests {
use super::*;
+ use crate::coalesce_partitions::CoalescePartitionsExec;
use crate::test::TestMemoryExec;
use crate::{
common, expressions::Column, repartition::RepartitionExec,
test::build_table_i32,
@@ -2105,6 +2101,7 @@ mod tests {
let left =
TestMemoryExec::try_new_exec(&[vec![batch1], vec![batch2]],
schema, None)
.unwrap();
+ let left = Arc::new(CoalescePartitionsExec::new(left));
let right = build_table(
("a1", &vec![1, 2, 3]),
@@ -2177,6 +2174,7 @@ mod tests {
let left =
TestMemoryExec::try_new_exec(&[vec![batch1], vec![batch2]],
schema, None)
.unwrap();
+ let left = Arc::new(CoalescePartitionsExec::new(left));
let right = build_table(
("a2", &vec![20, 30, 10]),
("b2", &vec![5, 6, 4]),
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
