feat: add LazyPartitioned mode for hash join to reduce RepartitionExec overhead

This commit adds a new PartitionMode::LazyPartitioned that avoids the
full build-side RepartitionExec when executing partitioned hash joins.
Instead of pre-repartitioning all columns of the build table, rows are
filtered lazily during hash table construction using hash(join_keys) %
partition_count.

Key changes:
- Add LazyPartitioned variant to PartitionMode enum
- Build side requests UnspecifiedDistribution (merged, no repartition)
- Probe side still requests HashPartitioned distribution
- Add filter_batch_by_partition() to filter build rows per partition
- Update collect_left_input to accept optional partition filter
- Add protobuf serialization support for new mode
- Update optimizer to handle LazyPartitioned in key reordering

This optimization is beneficial for wide build tables where copying
all columns in RepartitionExec is expensive.

Closes #19789

Author: GaneshPatil7517

datafusion/physical-optimizer/src/enforce_distribution.rs

@@ -327,6 +327,34 @@ pub fn adjust_input_keys_ordering(
                 )
                 .map(Transformed::yes);
             }
+            PartitionMode::LazyPartitioned => {
+                // LazyPartitioned mode uses the same key reordering as Partitioned,
+                // but with LazyPartitioned mode preserved
+                let join_constructor = |new_conditions: (
+                    Vec<(PhysicalExprRef, PhysicalExprRef)>,
+                    Vec<SortOptions>,
+                )| {
+                    HashJoinExec::try_new(
+                        Arc::clone(left),
+                        Arc::clone(right),
+                        new_conditions.0,
+                        filter.clone(),
+                        join_type,
+                        projection.clone(),
+                        PartitionMode::LazyPartitioned,
+                        *null_equality,
+                        *null_aware,
+                    )
+                    .map(|e| Arc::new(e) as _)
+                };
+                return reorder_partitioned_join_keys(
+                    requirements,
+                    on,
+                    &[],
+                    &join_constructor,
+                )
+                .map(Transformed::yes);
+            }
             PartitionMode::CollectLeft => {
                 // Push down requirements to the right side
                 requirements.children[1].data = match join_type {
@@ -624,7 +652,10 @@ pub fn reorder_join_keys_to_inputs(
         ..
     }) = plan_any.downcast_ref::<HashJoinExec>()
     {
-        if matches!(mode, PartitionMode::Partitioned) {
+        if matches!(
+            mode,
+            PartitionMode::Partitioned | PartitionMode::LazyPartitioned
+        ) {
             let (join_keys, positions) = reorder_current_join_keys(
                 extract_join_keys(on),
                 Some(left.output_partitioning()),
@@ -645,7 +676,7 @@ pub fn reorder_join_keys_to_inputs(
                     filter.clone(),
                     join_type,
                     projection.clone(),
-                    PartitionMode::Partitioned,
+                    *mode,
                     *null_equality,
                     *null_aware,
                 )?));
@@ -1257,6 +1288,10 @@ pub fn ensure_distribution(
     //
     // CollectLeft/CollectRight modes are safe because one side is collected
     // to a single partition which eliminates partition-to-partition mapping.
+    //
+    // LazyPartitioned mode is also safe from this issue because the build side
+    // is not pre-partitioned; instead, rows are filtered locally during hash
+    // table construction. Only the probe side is hash-partitioned.
     let is_partitioned_join = plan
         .as_any()
         .downcast_ref::<HashJoinExec>()

datafusion/physical-optimizer/src/join_selection.rs

@@ -293,7 +293,7 @@ fn statistical_join_selection_subrule(
                         || partitioned_hash_join(hash_join).map(Some),
                         |v| Ok(Some(v)),
                     )?,
-                PartitionMode::Partitioned => {
+                PartitionMode::Partitioned | PartitionMode::LazyPartitioned => {
                     let left = hash_join.left();
                     let right = hash_join.right();
                     // Don't swap null-aware anti joins as they have specific side requirements
@@ -302,7 +302,7 @@ fn statistical_join_selection_subrule(
                         && should_swap_join_order(&**left, &**right)?
                     {
                         hash_join
-                            .swap_inputs(PartitionMode::Partitioned)
+                            .swap_inputs(*hash_join.partition_mode())
                             .map(Some)?
                     } else {
                         None
@@ -540,6 +540,9 @@ pub(crate) fn swap_join_according_to_unboundedness(
         (PartitionMode::Partitioned, _) => {
             hash_join.swap_inputs(PartitionMode::Partitioned)
         }
+        (PartitionMode::LazyPartitioned, _) => {
+            hash_join.swap_inputs(PartitionMode::LazyPartitioned)
+        }
         (PartitionMode::CollectLeft, _) => {
             hash_join.swap_inputs(PartitionMode::CollectLeft)
         }

datafusion/physical-plan/src/joins/hash_join/exec.rs

@@ -27,6 +27,7 @@ use crate::filter_pushdown::{
     ChildPushdownResult, FilterDescription, FilterPushdownPhase,
     FilterPushdownPropagation,
 };
+use crate::hash_utils::create_hashes;
 use crate::joins::Map;
 use crate::joins::array_map::ArrayMap;
 use crate::joins::hash_join::inlist_builder::build_struct_inlist_values;
@@ -49,6 +50,7 @@ use crate::projection::{
 };
 use crate::repartition::REPARTITION_RANDOM_STATE;
 use crate::spill::get_record_batch_memory_size;
+use crate::stream::RecordBatchReceiverStream;
 use crate::{
     DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, Partitioning,
     PlanProperties, SendableRecordBatchStream, Statistics,
@@ -61,8 +63,8 @@ use crate::{
     metrics::{ExecutionPlanMetricsSet, MetricsSet},
 };
 
-use arrow::array::{ArrayRef, BooleanBufferBuilder};
-use arrow::compute::concat_batches;
+use arrow::array::{ArrayRef, BooleanArray, BooleanBufferBuilder};
+use arrow::compute::{concat_batches, filter_record_batch};
 use arrow::datatypes::SchemaRef;
 use arrow::record_batch::RecordBatch;
 use arrow::util::bit_util;
@@ -740,6 +742,12 @@ impl HashJoinExec {
             PartitionMode::Partitioned => {
                 symmetric_join_output_partitioning(left, right, &join_type)?
             }
+            PartitionMode::LazyPartitioned => {
+                // LazyPartitioned: output partitioning is determined by the probe (right) side,
+                // since each partition builds its own hash table from filtered build rows.
+                // This is similar to Partitioned mode but the build side isn't pre-partitioned.
+                symmetric_join_output_partitioning(left, right, &join_type)?
+            }
         };
 
         let emission_type = if left.boundedness().is_unbounded() {
@@ -958,6 +966,15 @@ impl ExecutionPlan for HashJoinExec {
                     Distribution::HashPartitioned(right_expr),
                 ]
             }
+            PartitionMode::LazyPartitioned => {
+                // LazyPartitioned mode: build side is NOT repartitioned (we read all
+                // partitions and filter locally), but probe side IS hash-partitioned.
+                let right_expr = self.on.iter().map(|(_, r)| Arc::clone(r)).collect();
+                vec![
+                    Distribution::UnspecifiedDistribution,
+                    Distribution::HashPartitioned(right_expr),
+                ]
+            }
             PartitionMode::Auto => vec![
                 Distribution::UnspecifiedDistribution,
                 Distribution::UnspecifiedDistribution,
@@ -1116,6 +1133,7 @@ impl ExecutionPlan for HashJoinExec {
                     Arc::clone(context.session_config().options()),
                     self.null_equality,
                     array_map_created_count,
+                    None, // No partition filtering for CollectLeft mode
                 ))
             })?,
             PartitionMode::Partitioned => {
@@ -1137,6 +1155,58 @@ impl ExecutionPlan for HashJoinExec {
                     Arc::clone(context.session_config().options()),
                     self.null_equality,
                     array_map_created_count,
+                    None, // No partition filtering - already pre-partitioned by RepartitionExec
+                ))
+            }
+            PartitionMode::LazyPartitioned => {
+                // LazyPartitioned mode: read ALL build partitions and filter locally
+                // by computing hash % num_partitions == current_partition.
+                // This avoids the overhead of RepartitionExec copying all columns.
+                let left_partitions = self.left.output_partitioning().partition_count();
+                let right_partitions = self.right.output_partitioning().partition_count();
+
+                let left_stream: SendableRecordBatchStream = if left_partitions == 1 {
+                    // Single input partition - use directly
+                    self.left.execute(0, Arc::clone(&context))?
+                } else {
+                    // Multiple input partitions - merge all into single stream
+                    let mut builder = RecordBatchReceiverStream::builder(
+                        self.left.schema(),
+                        left_partitions,
+                    );
+                    for part_i in 0..left_partitions {
+                        builder.run_input(
+                            Arc::clone(&self.left),
+                            part_i,
+                            Arc::clone(&context),
+                        );
+                    }
+                    builder.build()
+                };
+
+                let reservation =
+                    MemoryConsumer::new(format!("HashJoinInput[{partition}]"))
+                        .register(context.memory_pool());
+
+                // Create partition filter - filter rows to only keep those belonging to this partition
+                let partition_filter = Some(PartitionFilter {
+                    current_partition: partition,
+                    total_partitions: right_partitions,
+                });
+
+                OnceFut::new(collect_left_input(
+                    self.random_state.random_state().clone(),
+                    left_stream,
+                    on_left.clone(),
+                    join_metrics.clone(),
+                    reservation,
+                    need_produce_result_in_final(self.join_type),
+                    1, // Each partition has its own hash table
+                    enable_dynamic_filter_pushdown,
+                    Arc::clone(context.session_config().options()),
+                    self.null_equality,
+                    array_map_created_count,
+                    partition_filter,
                 ))
             }
             PartitionMode::Auto => {
@@ -1514,6 +1584,62 @@ fn should_collect_min_max_for_perfect_hash(
     Ok(ArrayMap::is_supported_type(&data_type))
 }
 
+/// Partition filter configuration for lazy partitioning.
+///
+/// When set, only rows where `hash(join_keys) % total_partitions == current_partition`
+/// are kept during hash table construction. This avoids the overhead of
+/// pre-partitioning with `RepartitionExec` by filtering rows locally.
+#[derive(Debug, Clone, Copy)]
+struct PartitionFilter {
+    /// The partition index this execution is responsible for
+    current_partition: usize,
+    /// Total number of partitions in the join
+    total_partitions: usize,
+}
+
+/// Filters a record batch to only include rows that belong to the specified partition.
+///
+/// Uses the same hash function and seeds as `RepartitionExec` to ensure that
+/// rows are routed consistently. Rows where `hash(join_keys) % total_partitions == current_partition`
+/// are kept; all others are filtered out.
+///
+/// This function is used in `LazyPartitioned` mode to avoid the overhead of
+/// `RepartitionExec` by filtering rows during hash table construction instead
+/// of pre-partitioning.
+fn filter_batch_by_partition(
+    batch: &RecordBatch,
+    on_left: &[PhysicalExprRef],
+    partition_filter: &PartitionFilter,
+) -> Result<RecordBatch> {
+    let num_rows = batch.num_rows();
+    if num_rows == 0 {
+        return Ok(batch.clone());
+    }
+
+    // Evaluate join key columns
+    let arrays = evaluate_expressions_to_arrays(on_left, batch)?;
+
+    // Compute hashes using the same random state as RepartitionExec
+    let mut hashes_buffer = vec![0u64; num_rows];
+    create_hashes(
+        &arrays,
+        REPARTITION_RANDOM_STATE.random_state(),
+        &mut hashes_buffer,
+    )?;
+
+    // Create a boolean mask for rows belonging to this partition
+    let mask: BooleanArray = hashes_buffer
+        .iter()
+        .map(|hash| {
+            *hash % partition_filter.total_partitions as u64
+                == partition_filter.current_partition as u64
+        })
+        .collect();
+
+    // Filter the batch
+    Ok(filter_record_batch(batch, &mask)?)
+}
+
 /// Collects all batches from the left (build) side stream and creates a hash map for joining.
 ///
 /// This function is responsible for:
@@ -1531,6 +1657,8 @@ fn should_collect_min_max_for_perfect_hash(
 /// * `with_visited_indices_bitmap` - Whether to track visited indices (for outer joins)
 /// * `probe_threads_count` - Number of threads that will probe this hash table
 /// * `should_compute_dynamic_filters` - Whether to compute min/max bounds for dynamic filtering
+/// * `partition_filter` - Optional partition filter for lazy partitioning mode.
+///   When set, only rows belonging to the specified partition are included in the hash table.
 ///
 /// # Dynamic Filter Coordination
 /// When `should_compute_dynamic_filters` is true, this function computes the min/max bounds
@@ -1555,6 +1683,7 @@ async fn collect_left_input(
     config: Arc<ConfigOptions>,
     null_equality: NullEquality,
     array_map_created_count: Count,
+    partition_filter: Option<PartitionFilter>,
 ) -> Result<JoinLeftData> {
     let schema = left_stream.schema();
 
@@ -1569,28 +1698,46 @@ async fn collect_left_input(
         should_compute_dynamic_filters || should_collect_min_max_for_phj,
     )?;
 
+    // Clone on_left for use in the closure
+    let on_left_for_filter = on_left.clone();
+
     let state = left_stream
-        .try_fold(initial, |mut state, batch| async move {
-            // Update accumulators if computing bounds
-            if let Some(ref mut accumulators) = state.bounds_accumulators {
-                for accumulator in accumulators {
-                    accumulator.update_batch(&batch)?;
+        .try_fold(initial, |mut state, batch| {
+            let on_left_clone = on_left_for_filter.clone();
+            async move {
+                // Apply partition filter if in lazy partitioning mode
+                let batch = if let Some(ref pf) = partition_filter {
+                    let filtered = filter_batch_by_partition(&batch, &on_left_clone, pf)?;
+                    // Skip empty batches after filtering
+                    if filtered.num_rows() == 0 {
+                        return Ok(state);
+                    }
+                    filtered
+                } else {
+                    batch
+                };
+
+                // Update accumulators if computing bounds
+                if let Some(ref mut accumulators) = state.bounds_accumulators {
+                    for accumulator in accumulators {
+                        accumulator.update_batch(&batch)?;
+                    }
                 }
-            }
 
-            // Decide if we spill or not
-            let batch_size = get_record_batch_memory_size(&batch);
-            // Reserve memory for incoming batch
-            state.reservation.try_grow(batch_size)?;
-            // Update metrics
-            state.metrics.build_mem_used.add(batch_size);
-            state.metrics.build_input_batches.add(1);
-            state.metrics.build_input_rows.add(batch.num_rows());
-            // Update row count
-            state.num_rows += batch.num_rows();
-            // Push batch to output
-            state.batches.push(batch);
-            Ok(state)
+                // Decide if we spill or not
+                let batch_size = get_record_batch_memory_size(&batch);
+                // Reserve memory for incoming batch
+                state.reservation.try_grow(batch_size)?;
+                // Update metrics
+                state.metrics.build_mem_used.add(batch_size);
+                state.metrics.build_input_batches.add(1);
+                state.metrics.build_input_rows.add(batch.num_rows());
+                // Update row count
+                state.num_rows += batch.num_rows();
+                // Push batch to output
+                state.batches.push(batch);
+                Ok(state)
+            }
         })
         .await?;
 
@@ -1983,6 +2130,10 @@ mod tests {
                 left,
                 Partitioning::Hash(left_expr, partition_count),
             )?),
+            PartitionMode::LazyPartitioned => {
+                // For LazyPartitioned, the build side is merged and filtered lazily per partition
+                Arc::new(CoalescePartitionsExec::new(left))
+            }
             PartitionMode::Auto => {
                 return internal_err!("Unexpected PartitionMode::Auto in join tests");
             }
@@ -2000,10 +2151,13 @@ mod tests {
                     Partitioning::Hash(partition_expr, partition_count),
                 )?) as _
             }
-            PartitionMode::Partitioned => Arc::new(RepartitionExec::try_new(
-                right,
-                Partitioning::Hash(right_expr, partition_count),
-            )?),
+            PartitionMode::Partitioned | PartitionMode::LazyPartitioned => {
+                // For both Partitioned and LazyPartitioned, probe side is hash partitioned
+                Arc::new(RepartitionExec::try_new(
+                    right,
+                    Partitioning::Hash(right_expr, partition_count),
+                )?)
+            }
             PartitionMode::Auto => {
                 return internal_err!("Unexpected PartitionMode::Auto in join tests");
             }