From 28b28bca0cc0b054bdf3521139ae6342f9cea723 Mon Sep 17 00:00:00 2001 From: LiaCastaneda Date: Tue, 14 Apr 2026 13:19:36 +0200 Subject: [PATCH 1/8] Make record_batch_buffer_size configurable --- src/distributed_planner/distributed_config.rs | 5 +++++ src/worker/impl_execute_task.rs | 5 ++--- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/src/distributed_planner/distributed_config.rs b/src/distributed_planner/distributed_config.rs index 8aeaa435..f4dc8284 100644 --- a/src/distributed_planner/distributed_config.rs +++ b/src/distributed_planner/distributed_config.rs @@ -48,6 +48,11 @@ extensions_options! { /// The compression used for sending data over the network between workers. /// It can be set to either `zstd`, `lz4` or `none`. pub compression: String, default = "lz4".to_string() + /// Size of the bounded mpsc channel used in `spawn_select_all` to merge N partition + /// streams into one gRPC response stream. With only 2 slots and N SpawnedTasks competing, + /// any brief gRPC/HTTP2 stall freezes the entire pipeline. Increase this experimentally + /// to measure the impact of the channel stall. Default is 2 to preserve existing behavior. + pub record_batch_buffer_size: usize, default = 2 /// Maximum tasks that will be assigned per stage during distributed planning. /// If set to 0, this value is the number of workers returned by the provided `WorkerResolver`. /// It defaults to 0. diff --git a/src/worker/impl_execute_task.rs b/src/worker/impl_execute_task.rs index e4eaa923..7ee7eea7 100644 --- a/src/worker/impl_execute_task.rs +++ b/src/worker/impl_execute_task.rs @@ -28,8 +28,6 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use tonic::{Request, Response, Status}; -/// How many record batches to buffer from the plan execution. -const RECORD_BATCH_BUFFER_SIZE: usize = 2; const WAIT_PLAN_TIMEOUT_SECS: u64 = 10; impl Worker { @@ -67,6 +65,7 @@ impl Worker { )))?, }; let send_metrics = d_cfg.collect_metrics; + let record_batch_buffer_size = d_cfg.record_batch_buffer_size; let partition_count = plan.properties().partitioning.partition_count(); let plan_name = plan.name(); @@ -160,7 +159,7 @@ impl Worker { // Merge all the per-partition streams into one. Each message in the stream is marked with // the original partition, so they can be reconstructed at the other side of the boundary. let memory_pool = Arc::clone(&task_ctx.runtime_env().memory_pool); - let stream = spawn_select_all(streams, memory_pool, RECORD_BATCH_BUFFER_SIZE); + let stream = spawn_select_all(streams, memory_pool, record_batch_buffer_size); Ok(Response::new(Box::pin(stream.map_err(|err| match err { FlightError::Tonic(status) => *status, From 66f56612a4dddf5e61d010274f17521c2fa895a0 Mon Sep 17 00:00:00 2001 From: Lavkesh Lahngir Date: Wed, 15 Apr 2026 16:51:55 +0200 Subject: [PATCH 2/8] network shuffle optimise --- src/distributed_planner/distributed_config.rs | 2 + .../distributed_physical_optimizer_rule.rs | 107 +++++++++++++- .../benchmarks/shuffle_bench.rs | 1 + src/execution_plans/network_shuffle.rs | 138 ++++++++++++++++-- src/protobuf/distributed_codec.rs | 1 + src/test_utils/plans.rs | 15 ++ 6 files changed, 250 insertions(+), 14 deletions(-) diff --git a/src/distributed_planner/distributed_config.rs b/src/distributed_planner/distributed_config.rs index f4dc8284..7698b695 100644 --- a/src/distributed_planner/distributed_config.rs +++ b/src/distributed_planner/distributed_config.rs @@ -45,6 +45,8 @@ extensions_options! { /// use broadcasting like checking build side size. /// For now, broadcasting all CollectLeft joins is not always beneficial. pub broadcast_joins: bool, default = false + /// Use 1:1 task-partition mapping instead of N*M fanout, eliminating partition inflation. + pub optimize_shuffle_partitioning: bool, default = false /// The compression used for sending data over the network between workers. /// It can be set to either `zstd`, `lz4` or `none`. pub compression: String, default = "lz4".to_string() diff --git a/src/distributed_planner/distributed_physical_optimizer_rule.rs b/src/distributed_planner/distributed_physical_optimizer_rule.rs index 59543c60..99102212 100644 --- a/src/distributed_planner/distributed_physical_optimizer_rule.rs +++ b/src/distributed_planner/distributed_physical_optimizer_rule.rs @@ -118,12 +118,22 @@ fn distribute_plan( if task_count == 1 && max_child_task_count == Some(1) { return require_one_child(new_children); } + + let child_plan = require_one_child(new_children)?; + + let downstream_task_count = if d_cfg.optimize_shuffle_partitioning { + child_plan.output_partitioning().partition_count() + } else { + task_count + }; + let node = Arc::new(NetworkShuffleExec::try_new( - require_one_child(new_children)?, + child_plan, query_id, *stage_id, - task_count, + downstream_task_count, max_child_task_count.unwrap_or(1), + d_cfg.optimize_shuffle_partitioning, )?); stage_id.add_assign(1); Ok(node) @@ -963,4 +973,97 @@ mod tests { format!("{}", displayable(physical_plan.as_ref()).indent(true)) } } + + async fn sql_to_explain_optimized( + query: &str, + num_workers: usize, + ) -> String { + use crate::test_utils::plans::base_session_builder_with_shuffle_optimization; + + let target_partitions = 4; + let mut builder = base_session_builder_with_shuffle_optimization( + target_partitions, + num_workers, + false, + true, + ); + + builder = builder.with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule)); + + let (ctx, query) = context_with_query(builder, query).await; + let df = ctx.sql(&query).await.unwrap(); + let physical_plan = df.create_physical_plan().await.unwrap(); + + display_plan_ascii(physical_plan.as_ref(), false) + } + + #[tokio::test] + async fn test_aggregation_optimized() { + let query = r#" + SELECT count(*), "RainToday" FROM weather GROUP BY "RainToday" ORDER BY count(*) + "#; + let plan = sql_to_explain_optimized(query, 3).await; + assert_snapshot!(plan, @r" + ┌───── DistributedExec ── Tasks: t0:[p0] + │ ProjectionExec: expr=[count(*)@0 as count(*), RainToday@1 as RainToday] + │ SortPreservingMergeExec: [count(Int64(1))@2 ASC NULLS LAST] + │ [Stage 2] => NetworkCoalesceExec: output_partitions=4, input_tasks=4 + └────────────────────────────────────────────────── + ┌───── Stage 2 ── Tasks: t0:[p0] t1:[p1] t2:[p2] t3:[p3] + │ SortExec: expr=[count(*)@0 ASC NULLS LAST], preserve_partitioning=[true] + │ ProjectionExec: expr=[count(Int64(1))@1 as count(*), RainToday@0 as RainToday, count(Int64(1))@1 as count(Int64(1))] + │ AggregateExec: mode=FinalPartitioned, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))] + │ [Stage 1] => NetworkShuffleExec: output_partitions=1, input_tasks=3 + └────────────────────────────────────────────────── + ┌───── Stage 1 ── Tasks: t0:[p0..p3] t1:[p0..p3] t2:[p0..p3] + │ RepartitionExec: partitioning=Hash([RainToday@0], 4), input_partitions=1 + │ AggregateExec: mode=Partial, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))] + │ PartitionIsolatorExec: t0:[p0,__,__] t1:[__,p0,__] t2:[__,__,p0] + │ DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet], [/testdata/weather/result-000001.parquet], [/testdata/weather/result-000002.parquet]]}, projection=[RainToday], file_type=parquet + └────────────────────────────────────────────────── + "); + } + + #[tokio::test] + async fn test_distinct_optimized() { + let query = r#" + SELECT DISTINCT "RainToday", "WindGustDir" FROM weather + "#; + let plan = sql_to_explain_optimized(query, 3).await; + assert_snapshot!(plan, @r" + ┌───── DistributedExec ── Tasks: t0:[p0] + │ CoalescePartitionsExec + │ [Stage 2] => NetworkCoalesceExec: output_partitions=4, input_tasks=4 + └────────────────────────────────────────────────── + ┌───── Stage 2 ── Tasks: t0:[p0] t1:[p1] t2:[p2] t3:[p3] + │ AggregateExec: mode=FinalPartitioned, gby=[RainToday@0 as RainToday, WindGustDir@1 as WindGustDir], aggr=[] + │ [Stage 1] => NetworkShuffleExec: output_partitions=1, input_tasks=3 + └────────────────────────────────────────────────── + ┌───── Stage 1 ── Tasks: t0:[p0..p3] t1:[p0..p3] t2:[p0..p3] + │ RepartitionExec: partitioning=Hash([RainToday@0, WindGustDir@1], 4), input_partitions=1 + │ AggregateExec: mode=Partial, gby=[RainToday@0 as RainToday, WindGustDir@1 as WindGustDir], aggr=[] + │ PartitionIsolatorExec: t0:[p0,__,__] t1:[__,p0,__] t2:[__,__,p0] + │ DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet], [/testdata/weather/result-000001.parquet], [/testdata/weather/result-000002.parquet]]}, projection=[RainToday, WindGustDir], file_type=parquet + └────────────────────────────────────────────────── + "); + } + + #[tokio::test] + async fn test_sort_optimized() { + let query = r#" + SELECT * FROM weather ORDER BY "MinTemp" DESC + "#; + let plan = sql_to_explain_optimized(query, 3).await; + assert_snapshot!(plan, @r" + ┌───── DistributedExec ── Tasks: t0:[p0] + │ SortPreservingMergeExec: [MinTemp@0 DESC] + │ [Stage 1] => NetworkCoalesceExec: output_partitions=3, input_tasks=3 + └────────────────────────────────────────────────── + ┌───── Stage 1 ── Tasks: t0:[p0] t1:[p1] t2:[p2] + │ SortExec: expr=[MinTemp@0 DESC], preserve_partitioning=[true] + │ PartitionIsolatorExec: t0:[p0,__,__] t1:[__,p0,__] t2:[__,__,p0] + │ DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet], [/testdata/weather/result-000001.parquet], [/testdata/weather/result-000002.parquet]]}, projection=[MinTemp, MaxTemp, Rainfall, Evaporation, Sunshine, WindGustDir, WindGustSpeed, WindDir9am, WindDir3pm, WindSpeed9am, WindSpeed3pm, Humidity9am, Humidity3pm, Pressure9am, Pressure3pm, Cloud9am, Cloud3pm, Temp9am, Temp3pm, RainToday, RISK_MM, RainTomorrow], file_type=parquet + └────────────────────────────────────────────────── + "); + } } diff --git a/src/execution_plans/benchmarks/shuffle_bench.rs b/src/execution_plans/benchmarks/shuffle_bench.rs index 51461dfa..214deaf6 100644 --- a/src/execution_plans/benchmarks/shuffle_bench.rs +++ b/src/execution_plans/benchmarks/shuffle_bench.rs @@ -149,6 +149,7 @@ impl ShuffleBench { input_stage: input_stage.clone(), worker_connections: WorkerConnectionPool::new(self.producer_tasks), metrics_collection: Arc::new(Default::default()), + optimize_shuffle_partitioning: false, }; let task_ctx = Arc::new(task_ctx_with_extension( &task_ctx, diff --git a/src/execution_plans/network_shuffle.rs b/src/execution_plans/network_shuffle.rs index 72557862..23880d5c 100644 --- a/src/execution_plans/network_shuffle.rs +++ b/src/execution_plans/network_shuffle.rs @@ -118,6 +118,8 @@ pub struct NetworkShuffleExec { /// a task to the last NetworkCoalesceExec to read from it, which may or may not be this /// instance. pub(crate) metrics_collection: Arc>>, + /// Whether to use optimized shuffle partitioning (1:1 task-partition mapping) + pub(crate) optimize_shuffle_partitioning: bool, } impl NetworkShuffleExec { @@ -131,6 +133,7 @@ impl NetworkShuffleExec { num: usize, task_count: usize, input_task_count: usize, + optimize_shuffle_partitioning: bool, ) -> Result { if !matches!(input.output_partitioning(), Partitioning::Hash(_, _)) { return plan_err!("NetworkShuffleExec input must be hash partitioned"); @@ -138,16 +141,19 @@ impl NetworkShuffleExec { let transformed = Arc::clone(&input).transform_down(|plan| { if let Some(r_exe) = plan.as_any().downcast_ref::() { - // Scale the input RepartitionExec to account for all the tasks to which it will - // need to fan data out. - let scaled = Arc::new(RepartitionExec::try_new( - require_one_child(r_exe.children())?, - scale_partitioning(r_exe.partitioning(), |p| p * task_count), - )?); + let scaled = if optimize_shuffle_partitioning { + Arc::new(RepartitionExec::try_new( + require_one_child(r_exe.children())?, + r_exe.partitioning().clone(), + )?) + } else { + Arc::new(RepartitionExec::try_new( + require_one_child(r_exe.children())?, + scale_partitioning(r_exe.partitioning(), |p| p * task_count), + )?) + }; Ok(Transformed::new(scaled, true, TreeNodeRecursion::Stop)) } else if matches!(plan.output_partitioning(), Partitioning::Hash(_, _)) { - // This might be a passthrough node, like a CoalesceBatchesExec or something like that. - // This is fine, we can let the node be here. Ok(Transformed::no(plan)) } else { plan_err!( @@ -157,6 +163,14 @@ impl NetworkShuffleExec { } })?; + let properties = if optimize_shuffle_partitioning { + let mut props = input.properties().clone(); + props.partitioning = Partitioning::UnknownPartitioning(1); + props + } else { + input.properties().clone() + }; + Ok(Self { input_stage: Stage { query_id, @@ -165,8 +179,9 @@ impl NetworkShuffleExec { tasks: vec![ExecutionTask { url: None }; input_task_count], }, worker_connections: WorkerConnectionPool::new(input_task_count), - properties: input.properties().clone(), + properties, metrics_collection: Default::default(), + optimize_shuffle_partitioning, }) } } @@ -230,19 +245,34 @@ impl ExecutionPlan for NetworkShuffleExec { context: Arc, ) -> Result { let task_context = DistributedTaskContext::from_ctx(&context); - let off = self.properties.partitioning.partition_count() * task_context.task_index; + + let (partition_offset, partition_range) = if self.optimize_shuffle_partitioning { + let partition_id = task_context.task_index; + (partition_id, partition_id..(partition_id + 1)) + } else { + let partition_count = self.properties.partitioning.partition_count(); + let off = partition_count * task_context.task_index; + (off, off..(off + partition_count)) + }; let mut streams = Vec::with_capacity(self.input_stage.tasks.len()); for input_task_index in 0..self.input_stage.tasks.len() { let worker_connection = self.worker_connections.get_or_init_worker_connection( &self.input_stage, - off..(off + self.properties.partitioning.partition_count()), + partition_range.clone(), input_task_index, &context, )?; let metrics_collection = Arc::clone(&self.metrics_collection); - let stream = worker_connection.stream_partition(off + partition, move |meta| { + + let actual_partition = if self.optimize_shuffle_partitioning { + partition_offset + } else { + partition_offset + partition + }; + + let stream = worker_connection.stream_partition(actual_partition, move |meta| { if let Some(flight_app_metadata::Content::MetricsCollection(m)) = meta.content { for task_metrics in m.tasks { if let Some(task_key) = task_metrics.task_key { @@ -264,3 +294,87 @@ impl ExecutionPlan for NetworkShuffleExec { Some(self.worker_connections.metrics.clone_inner()) } } + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::physical_expr::expressions::Column; + use datafusion::physical_plan::empty::EmptyExec; + use datafusion::physical_plan::repartition::RepartitionExec; + use uuid::Uuid; + + #[test] + fn test_optimized_mode_sets_single_partition() -> datafusion::common::Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let part = Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 8); + + let repartition = Arc::new(RepartitionExec::try_new( + Arc::new(EmptyExec::new(schema.clone())), + part.clone(), + )?); + + let optimized = NetworkShuffleExec::try_new( + repartition.clone(), + Uuid::new_v4(), + 1, + 8, + 3, + true, + )?; + + assert_eq!(optimized.properties.partitioning.partition_count(), 1); + assert!(optimized.optimize_shuffle_partitioning); + + let legacy = NetworkShuffleExec::try_new( + repartition.clone(), + Uuid::new_v4(), + 1, + 2, + 3, + false, + )?; + + assert_eq!(legacy.properties.partitioning.partition_count(), 8); + assert!(!legacy.optimize_shuffle_partitioning); + + Ok(()) + } + + #[test] + fn test_optimized_mode_no_partition_inflation() -> datafusion::common::Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("b", DataType::Int32, false)])); + let part = Partitioning::Hash(vec![Arc::new(Column::new("b", 0))], 4); + + let repartition = Arc::new(RepartitionExec::try_new( + Arc::new(EmptyExec::new(schema.clone())), + part.clone(), + )?); + + let optimized = NetworkShuffleExec::try_new( + repartition.clone(), + Uuid::new_v4(), + 1, + 4, + 2, + true, + )?; + + let input_plan = optimized.input_stage.plan.as_ref().unwrap(); + assert_eq!(input_plan.output_partitioning().partition_count(), 4); + + let legacy = NetworkShuffleExec::try_new( + repartition, + Uuid::new_v4(), + 1, + 4, + 2, + false, + )?; + + let input_plan = legacy.input_stage.plan.as_ref().unwrap(); + assert_eq!(input_plan.output_partitioning().partition_count(), 16); + + Ok(()) + } +} diff --git a/src/protobuf/distributed_codec.rs b/src/protobuf/distributed_codec.rs index c81e3327..5842396e 100644 --- a/src/protobuf/distributed_codec.rs +++ b/src/protobuf/distributed_codec.rs @@ -437,6 +437,7 @@ fn new_network_hash_shuffle_exec( worker_connections: WorkerConnectionPool::new(input_stage.tasks.len()), input_stage, metrics_collection: Default::default(), + optimize_shuffle_partitioning: false, } } diff --git a/src/test_utils/plans.rs b/src/test_utils/plans.rs index 635bed47..9e372fd9 100644 --- a/src/test_utils/plans.rs +++ b/src/test_utils/plans.rs @@ -117,6 +117,20 @@ pub(crate) fn base_session_builder( target_partitions: usize, num_workers: usize, broadcast_enabled: bool, +) -> SessionStateBuilder { + base_session_builder_with_shuffle_optimization( + target_partitions, + num_workers, + broadcast_enabled, + false, + ) +} + +pub(crate) fn base_session_builder_with_shuffle_optimization( + target_partitions: usize, + num_workers: usize, + broadcast_enabled: bool, + optimize_shuffle_partitioning: bool, ) -> SessionStateBuilder { let mut config = SessionConfig::new() .with_target_partitions(target_partitions) @@ -124,6 +138,7 @@ pub(crate) fn base_session_builder( let d_cfg = DistributedConfig { broadcast_joins: broadcast_enabled, + optimize_shuffle_partitioning, ..Default::default() }; config.set_distributed_option_extension(d_cfg); From c23d51857809dddc60922236c49a19ee5a4c85a9 Mon Sep 17 00:00:00 2001 From: Lavkesh Lahngir Date: Thu, 16 Apr 2026 17:23:18 +0200 Subject: [PATCH 3/8] optimize_shuffle_partitioning in networkshuffle proto --- .../distributed_physical_optimizer_rule.rs | 4 ++-- src/execution_plans/network_shuffle.rs | 15 +++++++++++---- src/protobuf/distributed_codec.rs | 16 ++++++++++++++-- 3 files changed, 27 insertions(+), 8 deletions(-) diff --git a/src/distributed_planner/distributed_physical_optimizer_rule.rs b/src/distributed_planner/distributed_physical_optimizer_rule.rs index 99102212..76951bf6 100644 --- a/src/distributed_planner/distributed_physical_optimizer_rule.rs +++ b/src/distributed_planner/distributed_physical_optimizer_rule.rs @@ -1013,7 +1013,7 @@ mod tests { │ SortExec: expr=[count(*)@0 ASC NULLS LAST], preserve_partitioning=[true] │ ProjectionExec: expr=[count(Int64(1))@1 as count(*), RainToday@0 as RainToday, count(Int64(1))@1 as count(Int64(1))] │ AggregateExec: mode=FinalPartitioned, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))] - │ [Stage 1] => NetworkShuffleExec: output_partitions=1, input_tasks=3 + │ [Stage 1] => NetworkShuffleExec: output_partitions=1, input_tasks=3, optimized=true └────────────────────────────────────────────────── ┌───── Stage 1 ── Tasks: t0:[p0..p3] t1:[p0..p3] t2:[p0..p3] │ RepartitionExec: partitioning=Hash([RainToday@0], 4), input_partitions=1 @@ -1037,7 +1037,7 @@ mod tests { └────────────────────────────────────────────────── ┌───── Stage 2 ── Tasks: t0:[p0] t1:[p1] t2:[p2] t3:[p3] │ AggregateExec: mode=FinalPartitioned, gby=[RainToday@0 as RainToday, WindGustDir@1 as WindGustDir], aggr=[] - │ [Stage 1] => NetworkShuffleExec: output_partitions=1, input_tasks=3 + │ [Stage 1] => NetworkShuffleExec: output_partitions=1, input_tasks=3, optimized=true └────────────────────────────────────────────────── ┌───── Stage 1 ── Tasks: t0:[p0..p3] t1:[p0..p3] t2:[p0..p3] │ RepartitionExec: partitioning=Hash([RainToday@0, WindGustDir@1], 4), input_partitions=1 diff --git a/src/execution_plans/network_shuffle.rs b/src/execution_plans/network_shuffle.rs index 23880d5c..7c349d39 100644 --- a/src/execution_plans/network_shuffle.rs +++ b/src/execution_plans/network_shuffle.rs @@ -203,10 +203,17 @@ impl DisplayAs for NetworkShuffleExec { let input_tasks = self.input_stage.tasks.len(); let partitions = self.properties.partitioning.partition_count(); let stage = self.input_stage.num; - write!( - f, - "[Stage {stage}] => NetworkShuffleExec: output_partitions={partitions}, input_tasks={input_tasks}", - ) + if self.optimize_shuffle_partitioning { + write!( + f, + "[Stage {stage}] => NetworkShuffleExec: output_partitions={partitions}, input_tasks={input_tasks}, optimized=true", + ) + } else { + write!( + f, + "[Stage {stage}] => NetworkShuffleExec: output_partitions={partitions}, input_tasks={input_tasks}", + ) + } } } diff --git a/src/protobuf/distributed_codec.rs b/src/protobuf/distributed_codec.rs index 5842396e..c25f8977 100644 --- a/src/protobuf/distributed_codec.rs +++ b/src/protobuf/distributed_codec.rs @@ -80,6 +80,7 @@ impl PhysicalExtensionCodec for DistributedCodec { schema, partitioning, input_stage, + optimize_shuffle_partitioning, }) => { let schema: Schema = schema .as_ref() @@ -99,6 +100,7 @@ impl PhysicalExtensionCodec for DistributedCodec { partitioning, Arc::new(schema), parse_stage_proto(input_stage, inputs)?, + optimize_shuffle_partitioning, ))) } DistributedExecNode::NetworkCoalesceTasks(NetworkCoalesceExecProto { @@ -246,6 +248,7 @@ impl PhysicalExtensionCodec for DistributedCodec { &DefaultPhysicalProtoConverter, )?), input_stage: Some(encode_stage_proto(node.input_stage())?), + optimize_shuffle_partitioning: node.optimize_shuffle_partitioning, }; let wrapper = DistributedExecProto { @@ -396,6 +399,8 @@ pub struct NetworkShuffleExecProto { partitioning: Option, #[prost(message, optional, tag = "3")] input_stage: Option, + #[prost(bool, tag = "4")] + optimize_shuffle_partitioning: bool, } #[derive(Clone, PartialEq, ::prost::Message)] @@ -426,6 +431,7 @@ fn new_network_hash_shuffle_exec( partitioning: Partitioning, schema: SchemaRef, input_stage: Stage, + optimize_shuffle_partitioning: bool, ) -> NetworkShuffleExec { NetworkShuffleExec { properties: Arc::new(PlanProperties::new( @@ -437,7 +443,7 @@ fn new_network_hash_shuffle_exec( worker_connections: WorkerConnectionPool::new(input_stage.tasks.len()), input_stage, metrics_collection: Default::default(), - optimize_shuffle_partitioning: false, + optimize_shuffle_partitioning, } } @@ -586,7 +592,7 @@ mod tests { let schema = schema_i32("a"); let part = Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 4); let plan: Arc = - Arc::new(new_network_hash_shuffle_exec(part, schema, dummy_stage())); + Arc::new(new_network_hash_shuffle_exec(part, schema, dummy_stage(), false)); let mut buf = Vec::new(); codec.try_encode(plan.clone(), &mut buf)?; @@ -607,6 +613,7 @@ mod tests { Partitioning::UnknownPartitioning(1), schema, dummy_stage(), + false, )); let plan: Arc = Arc::new(PartitionIsolatorExec::new(flight.clone(), 1)); @@ -630,11 +637,13 @@ mod tests { Partitioning::RoundRobinBatch(2), schema.clone(), dummy_stage(), + false, )); let right = Arc::new(new_network_hash_shuffle_exec( Partitioning::RoundRobinBatch(2), schema.clone(), dummy_stage(), + false, )); let union = UnionExec::try_new(vec![left.clone(), right.clone()])?; @@ -659,6 +668,7 @@ mod tests { Partitioning::UnknownPartitioning(1), schema.clone(), dummy_stage(), + false, )); let sort_expr = PhysicalSortExpr { @@ -807,11 +817,13 @@ mod tests { Partitioning::RoundRobinBatch(2), schema.clone(), dummy_stage(), + false, )) as Arc; let right = Arc::new(new_network_hash_shuffle_exec( Partitioning::RoundRobinBatch(2), schema.clone(), dummy_stage(), + false, )) as Arc; let plan: Arc = From d181f3c49bd3f2256c323debb5a3e6d7ca8e8c7d Mon Sep 17 00:00:00 2001 From: Lavkesh Lahngir Date: Thu, 16 Apr 2026 21:31:19 +0200 Subject: [PATCH 4/8] add optimized flag value in the plan --- src/distributed_ext.rs | 41 +++++++++++++++++++ ...tch_coalescing_below_network_boundaries.rs | 6 +-- .../distributed_physical_optimizer_rule.rs | 22 ++++------ src/distributed_planner/plan_annotator.rs | 17 ++++++++ src/execution_plans/network_shuffle.rs | 16 +++----- 5 files changed, 74 insertions(+), 28 deletions(-) diff --git a/src/distributed_ext.rs b/src/distributed_ext.rs index 7c957926..5a73f040 100644 --- a/src/distributed_ext.rs +++ b/src/distributed_ext.rs @@ -455,6 +455,19 @@ pub trait DistributedExt: Sized { /// Same as [DistributedExt::with_distributed_broadcast_joins_enabled] but with an in-place mutation. fn set_distributed_broadcast_joins(&mut self, enabled: bool) -> Result<(), DataFusionError>; + /// Enables or disables optimized shuffle partitioning (1:1 task-partition mapping). + /// When enabled, eliminates N×M partition inflation during network shuffles. + fn with_distributed_optimize_shuffle_partitioning( + self, + enabled: bool, + ) -> Result; + + /// Same as [DistributedExt::with_distributed_optimize_shuffle_partitioning] but with an in-place mutation. + fn set_distributed_optimize_shuffle_partitioning( + &mut self, + enabled: bool, + ) -> Result<(), DataFusionError>; + /// The compression type to use for sending data over the wire. /// /// The default is [CompressionType::LZ4_FRAME]. @@ -615,6 +628,15 @@ impl DistributedExt for SessionConfig { Ok(()) } + fn set_distributed_optimize_shuffle_partitioning( + &mut self, + enabled: bool, + ) -> Result<(), DataFusionError> { + let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?; + d_cfg.optimize_shuffle_partitioning = enabled; + Ok(()) + } + fn set_distributed_compression( &mut self, compression: Option, @@ -703,6 +725,10 @@ impl DistributedExt for SessionConfig { #[expr($?;Ok(self))] fn with_distributed_broadcast_joins(mut self, enabled: bool) -> Result; + #[call(set_distributed_optimize_shuffle_partitioning)] + #[expr($?;Ok(self))] + fn with_distributed_optimize_shuffle_partitioning(mut self, enabled: bool) -> Result; + #[call(set_distributed_compression)] #[expr($?;Ok(self))] fn with_distributed_compression(mut self, compression: Option) -> Result; @@ -785,6 +811,11 @@ impl DistributedExt for SessionStateBuilder { #[expr($?;Ok(self))] fn with_distributed_broadcast_joins(mut self, enabled: bool) -> Result; + fn set_distributed_optimize_shuffle_partitioning(&mut self, enabled: bool) -> Result<(), DataFusionError>; + #[call(set_distributed_optimize_shuffle_partitioning)] + #[expr($?;Ok(self))] + fn with_distributed_optimize_shuffle_partitioning(mut self, enabled: bool) -> Result; + fn set_distributed_compression(&mut self, compression: Option) -> Result<(), DataFusionError>; #[call(set_distributed_compression)] #[expr($?;Ok(self))] @@ -871,6 +902,11 @@ impl DistributedExt for SessionState { #[expr($?;Ok(self))] fn with_distributed_broadcast_joins(mut self, enabled: bool) -> Result; + fn set_distributed_optimize_shuffle_partitioning(&mut self, enabled: bool) -> Result<(), DataFusionError>; + #[call(set_distributed_optimize_shuffle_partitioning)] + #[expr($?;Ok(self))] + fn with_distributed_optimize_shuffle_partitioning(mut self, enabled: bool) -> Result; + fn set_distributed_compression(&mut self, compression: Option) -> Result<(), DataFusionError>; #[call(set_distributed_compression)] #[expr($?;Ok(self))] @@ -957,6 +993,11 @@ impl DistributedExt for SessionContext { #[expr($?;Ok(self))] fn with_distributed_broadcast_joins(self, enabled: bool) -> Result; + fn set_distributed_optimize_shuffle_partitioning(&mut self, enabled: bool) -> Result<(), DataFusionError>; + #[call(set_distributed_optimize_shuffle_partitioning)] + #[expr($?;Ok(self))] + fn with_distributed_optimize_shuffle_partitioning(self, enabled: bool) -> Result; + fn set_distributed_compression(&mut self, compression: Option) -> Result<(), DataFusionError>; #[call(set_distributed_compression)] #[expr($?;Ok(self))] diff --git a/src/distributed_planner/batch_coalescing_below_network_boundaries.rs b/src/distributed_planner/batch_coalescing_below_network_boundaries.rs index f84443fa..eaac8d30 100644 --- a/src/distributed_planner/batch_coalescing_below_network_boundaries.rs +++ b/src/distributed_planner/batch_coalescing_below_network_boundaries.rs @@ -87,7 +87,7 @@ mod tests { └────────────────────────────────────────────────── ┌───── Stage 2 ── Tasks: t0:[p0..p3] t1:[p0..p3] │ AggregateExec: mode=FinalPartitioned, gby=[RainToday@0 as RainToday, WindGustDir@1 as WindGustDir], aggr=[] - │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=3 + │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=3, optimized=false └────────────────────────────────────────────────── ┌───── Stage 1 ── Tasks: t0:[p0..p7] t1:[p0..p7] t2:[p0..p7] │ RepartitionExec: partitioning=Hash([RainToday@0, WindGustDir@1], 8), input_partitions=4 @@ -115,7 +115,7 @@ mod tests { └────────────────────────────────────────────────── ┌───── Stage 2 ── Tasks: t0:[p0..p3] t1:[p0..p3] │ AggregateExec: mode=FinalPartitioned, gby=[RainToday@0 as RainToday, WindGustDir@1 as WindGustDir], aggr=[] - │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=3 + │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=3, optimized=false └────────────────────────────────────────────────── ┌───── Stage 1 ── Tasks: t0:[p0..p7] t1:[p0..p7] t2:[p0..p7] │ RepartitionExec: partitioning=Hash([RainToday@0, WindGustDir@1], 8), input_partitions=4 @@ -144,7 +144,7 @@ mod tests { ┌───── Stage 2 ── Tasks: t0:[p0..p3] t1:[p0..p3] │ CoalesceBatchesExec: target_batch_size=101 │ AggregateExec: mode=FinalPartitioned, gby=[RainToday@0 as RainToday, WindGustDir@1 as WindGustDir], aggr=[] - │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=3 + │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=3, optimized=false └────────────────────────────────────────────────── ┌───── Stage 1 ── Tasks: t0:[p0..p7] t1:[p0..p7] t2:[p0..p7] │ CoalesceBatchesExec: target_batch_size=101 diff --git a/src/distributed_planner/distributed_physical_optimizer_rule.rs b/src/distributed_planner/distributed_physical_optimizer_rule.rs index 76951bf6..0392c530 100644 --- a/src/distributed_planner/distributed_physical_optimizer_rule.rs +++ b/src/distributed_planner/distributed_physical_optimizer_rule.rs @@ -121,17 +121,11 @@ fn distribute_plan( let child_plan = require_one_child(new_children)?; - let downstream_task_count = if d_cfg.optimize_shuffle_partitioning { - child_plan.output_partitioning().partition_count() - } else { - task_count - }; - let node = Arc::new(NetworkShuffleExec::try_new( child_plan, query_id, *stage_id, - downstream_task_count, + task_count, max_child_task_count.unwrap_or(1), d_cfg.optimize_shuffle_partitioning, )?); @@ -247,7 +241,7 @@ mod tests { │ SortExec: expr=[count(*)@0 ASC NULLS LAST], preserve_partitioning=[true] │ ProjectionExec: expr=[count(Int64(1))@1 as count(*), RainToday@0 as RainToday, count(Int64(1))@1 as count(Int64(1))] │ AggregateExec: mode=FinalPartitioned, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))] - │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=3 + │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=3, optimized=false └────────────────────────────────────────────────── ┌───── Stage 1 ── Tasks: t0:[p0..p7] t1:[p0..p7] t2:[p0..p7] │ RepartitionExec: partitioning=Hash([RainToday@0], 8), input_partitions=1 @@ -277,7 +271,7 @@ mod tests { │ SortExec: expr=[count(*)@0 ASC NULLS LAST], preserve_partitioning=[true] │ ProjectionExec: expr=[count(Int64(1))@1 as count(*), RainToday@0 as RainToday, count(Int64(1))@1 as count(Int64(1))] │ AggregateExec: mode=FinalPartitioned, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))] - │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=2 + │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=2, optimized=false └────────────────────────────────────────────────── ┌───── Stage 1 ── Tasks: t0:[p0..p7] t1:[p0..p7] │ RepartitionExec: partitioning=Hash([RainToday@0], 8), input_partitions=2 @@ -327,7 +321,7 @@ mod tests { │ SortExec: expr=[count(*)@0 ASC NULLS LAST], preserve_partitioning=[true] │ ProjectionExec: expr=[count(Int64(1))@1 as count(*), RainToday@0 as RainToday, count(Int64(1))@1 as count(Int64(1))] │ AggregateExec: mode=FinalPartitioned, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))] - │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=3 + │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=3, optimized=false └────────────────────────────────────────────────── ┌───── Stage 1 ── Tasks: t0:[p0..p3] t1:[p0..p3] t2:[p0..p3] │ RepartitionExec: partitioning=Hash([RainToday@0], 4), input_partitions=1 @@ -380,7 +374,7 @@ mod tests { │ SortExec: expr=[count(*)@0 ASC NULLS LAST], preserve_partitioning=[true] │ ProjectionExec: expr=[count(Int64(1))@1 as count(*), RainToday@0 as RainToday, count(Int64(1))@1 as count(Int64(1))] │ AggregateExec: mode=FinalPartitioned, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))] - │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=3 + │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=3, optimized=false └────────────────────────────────────────────────── ┌───── Stage 1 ── Tasks: t0:[p0..p7] t1:[p0..p7] t2:[p0..p7] │ RepartitionExec: partitioning=Hash([RainToday@0], 8), input_partitions=1 @@ -445,12 +439,12 @@ mod tests { │ [Stage 2] => NetworkCoalesceExec: output_partitions=8, input_tasks=2 │ ProjectionExec: expr=[avg(weather.MaxTemp)@1 as MaxTemp, RainTomorrow@0 as RainTomorrow] │ AggregateExec: mode=FinalPartitioned, gby=[RainTomorrow@0 as RainTomorrow], aggr=[avg(weather.MaxTemp)] - │ [Stage 3] => NetworkShuffleExec: output_partitions=4, input_tasks=3 + │ [Stage 3] => NetworkShuffleExec: output_partitions=4, input_tasks=3, optimized=false └────────────────────────────────────────────────── ┌───── Stage 2 ── Tasks: t0:[p0..p3] t1:[p0..p3] │ ProjectionExec: expr=[avg(weather.MinTemp)@1 as MinTemp, RainTomorrow@0 as RainTomorrow] │ AggregateExec: mode=FinalPartitioned, gby=[RainTomorrow@0 as RainTomorrow], aggr=[avg(weather.MinTemp)] - │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=3 + │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=3, optimized=false └────────────────────────────────────────────────── ┌───── Stage 1 ── Tasks: t0:[p0..p7] t1:[p0..p7] t2:[p0..p7] │ RepartitionExec: partitioning=Hash([RainTomorrow@0], 8), input_partitions=4 @@ -509,7 +503,7 @@ mod tests { └────────────────────────────────────────────────── ┌───── Stage 2 ── Tasks: t0:[p0..p3] t1:[p0..p3] │ AggregateExec: mode=FinalPartitioned, gby=[RainToday@0 as RainToday, WindGustDir@1 as WindGustDir], aggr=[] - │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=3 + │ [Stage 1] => NetworkShuffleExec: output_partitions=4, input_tasks=3, optimized=false └────────────────────────────────────────────────── ┌───── Stage 1 ── Tasks: t0:[p0..p7] t1:[p0..p7] t2:[p0..p7] │ RepartitionExec: partitioning=Hash([RainToday@0, WindGustDir@1], 8), input_partitions=1 diff --git a/src/distributed_planner/plan_annotator.rs b/src/distributed_planner/plan_annotator.rs index 9aeeb461..7fbd0d19 100644 --- a/src/distributed_planner/plan_annotator.rs +++ b/src/distributed_planner/plan_annotator.rs @@ -5,6 +5,7 @@ use datafusion::common::{DataFusionError, plan_datafusion_err}; use datafusion::config::ConfigOptions; use datafusion::physical_expr::Partitioning; use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::ExecutionPlanProperties; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion::physical_plan::execution_plan::CardinalityEffect; use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode}; @@ -386,6 +387,22 @@ fn _annotate_plan( ); let prev_task_count = annotation.task_count.as_usize() as f64; annotation.task_count = Desired((prev_task_count * sf).ceil() as usize); + + // In optimized shuffle mode, NetworkShuffleExec uses task_index directly as the partition + // ID to fetch from the upstream stage. This requires downstream task count == upstream + // partition count. Override the cardinality-derived downstream task count here to enforce + // that invariant so stage scheduling spawns the right number of tasks. + if d_cfg.optimize_shuffle_partitioning { + if let PlanOrNetworkBoundary::Shuffle = &annotation.plan_or_nb { + if let Some(PlanOrNetworkBoundary::Plan(repartition)) = + annotation.children.first().map(|c| &c.plan_or_nb) + { + let partition_count = repartition.output_partitioning().partition_count(); + annotation.task_count = Desired(partition_count); + } + } + } + Ok(annotation) } else if root { // If this is the root node, it means that we have just finished annotating nodes for the diff --git a/src/execution_plans/network_shuffle.rs b/src/execution_plans/network_shuffle.rs index 7c349d39..a2b15782 100644 --- a/src/execution_plans/network_shuffle.rs +++ b/src/execution_plans/network_shuffle.rs @@ -203,17 +203,11 @@ impl DisplayAs for NetworkShuffleExec { let input_tasks = self.input_stage.tasks.len(); let partitions = self.properties.partitioning.partition_count(); let stage = self.input_stage.num; - if self.optimize_shuffle_partitioning { - write!( - f, - "[Stage {stage}] => NetworkShuffleExec: output_partitions={partitions}, input_tasks={input_tasks}, optimized=true", - ) - } else { - write!( - f, - "[Stage {stage}] => NetworkShuffleExec: output_partitions={partitions}, input_tasks={input_tasks}", - ) - } + let optimized = self.optimize_shuffle_partitioning; + write!( + f, + "[Stage {stage}] => NetworkShuffleExec: output_partitions={partitions}, input_tasks={input_tasks}, optimized={optimized}", + ) } } From b185751b619698a085d389c62a636775227e6a8d Mon Sep 17 00:00:00 2001 From: Lavkesh Lahngir Date: Fri, 17 Apr 2026 01:35:23 +0200 Subject: [PATCH 5/8] per boundary --- .../distributed_physical_optimizer_rule.rs | 9 ++++++++- src/distributed_planner/plan_annotator.rs | 8 +++++++- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/src/distributed_planner/distributed_physical_optimizer_rule.rs b/src/distributed_planner/distributed_physical_optimizer_rule.rs index 0392c530..e8950b64 100644 --- a/src/distributed_planner/distributed_physical_optimizer_rule.rs +++ b/src/distributed_planner/distributed_physical_optimizer_rule.rs @@ -121,13 +121,20 @@ fn distribute_plan( let child_plan = require_one_child(new_children)?; + // Only use optimized shuffle when the annotated task_count matches the upstream + // partition count exactly — meaning the annotator determined the optimization was + // beneficial (cardinality estimate <= partition count). Otherwise fall back to + // non-optimized to preserve downstream parallelism. + let use_optimized = d_cfg.optimize_shuffle_partitioning + && task_count == child_plan.output_partitioning().partition_count(); + let node = Arc::new(NetworkShuffleExec::try_new( child_plan, query_id, *stage_id, task_count, max_child_task_count.unwrap_or(1), - d_cfg.optimize_shuffle_partitioning, + use_optimized, )?); stage_id.add_assign(1); Ok(node) diff --git a/src/distributed_planner/plan_annotator.rs b/src/distributed_planner/plan_annotator.rs index 7fbd0d19..7b5338d7 100644 --- a/src/distributed_planner/plan_annotator.rs +++ b/src/distributed_planner/plan_annotator.rs @@ -392,13 +392,19 @@ fn _annotate_plan( // ID to fetch from the upstream stage. This requires downstream task count == upstream // partition count. Override the cardinality-derived downstream task count here to enforce // that invariant so stage scheduling spawns the right number of tasks. + // + // Only apply when the cardinality-derived task count does not exceed the upstream partition + // count. If cardinality says more tasks are needed (high-volume stages), fall back to + // non-optimized so parallelism is preserved. if d_cfg.optimize_shuffle_partitioning { if let PlanOrNetworkBoundary::Shuffle = &annotation.plan_or_nb { if let Some(PlanOrNetworkBoundary::Plan(repartition)) = annotation.children.first().map(|c| &c.plan_or_nb) { let partition_count = repartition.output_partitioning().partition_count(); - annotation.task_count = Desired(partition_count); + if annotation.task_count.as_usize() <= partition_count { + annotation.task_count = Desired(partition_count); + } } } } From 8bead7d1e12ec6e999ca671e9e8452ad1d3dd29f Mon Sep 17 00:00:00 2001 From: Lavkesh Lahngir Date: Fri, 17 Apr 2026 02:13:58 +0200 Subject: [PATCH 6/8] add threshold for optimizer --- src/distributed_ext.rs | 45 +++++++++++++++++++ src/distributed_planner/distributed_config.rs | 5 +++ src/distributed_planner/plan_annotator.rs | 3 +- 3 files changed, 52 insertions(+), 1 deletion(-) diff --git a/src/distributed_ext.rs b/src/distributed_ext.rs index 5a73f040..8a6c8a1b 100644 --- a/src/distributed_ext.rs +++ b/src/distributed_ext.rs @@ -468,6 +468,19 @@ pub trait DistributedExt: Sized { enabled: bool, ) -> Result<(), DataFusionError>; + /// Scale factor controlling when optimized shuffle partitioning is applied at a boundary. + /// The optimization is applied when cardinality_task_count <= partition_count * ratio. + fn with_distributed_optimize_shuffle_partitioning_ratio( + self, + ratio: f64, + ) -> Result; + + /// Same as [DistributedExt::with_distributed_optimize_shuffle_partitioning_ratio] but with an in-place mutation. + fn set_distributed_optimize_shuffle_partitioning_ratio( + &mut self, + ratio: f64, + ) -> Result<(), DataFusionError>; + /// The compression type to use for sending data over the wire. /// /// The default is [CompressionType::LZ4_FRAME]. @@ -637,6 +650,23 @@ impl DistributedExt for SessionConfig { Ok(()) } + fn set_distributed_optimize_shuffle_partitioning_ratio( + &mut self, + ratio: f64, + ) -> Result<(), DataFusionError> { + let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?; + d_cfg.optimize_shuffle_partitioning_ratio = ratio; + Ok(()) + } + + fn with_distributed_optimize_shuffle_partitioning_ratio( + mut self, + ratio: f64, + ) -> Result { + self.set_distributed_optimize_shuffle_partitioning_ratio(ratio)?; + Ok(self) + } + fn set_distributed_compression( &mut self, compression: Option, @@ -816,6 +846,11 @@ impl DistributedExt for SessionStateBuilder { #[expr($?;Ok(self))] fn with_distributed_optimize_shuffle_partitioning(mut self, enabled: bool) -> Result; + fn set_distributed_optimize_shuffle_partitioning_ratio(&mut self, ratio: f64) -> Result<(), DataFusionError>; + #[call(set_distributed_optimize_shuffle_partitioning_ratio)] + #[expr($?;Ok(self))] + fn with_distributed_optimize_shuffle_partitioning_ratio(mut self, ratio: f64) -> Result; + fn set_distributed_compression(&mut self, compression: Option) -> Result<(), DataFusionError>; #[call(set_distributed_compression)] #[expr($?;Ok(self))] @@ -907,6 +942,11 @@ impl DistributedExt for SessionState { #[expr($?;Ok(self))] fn with_distributed_optimize_shuffle_partitioning(mut self, enabled: bool) -> Result; + fn set_distributed_optimize_shuffle_partitioning_ratio(&mut self, ratio: f64) -> Result<(), DataFusionError>; + #[call(set_distributed_optimize_shuffle_partitioning_ratio)] + #[expr($?;Ok(self))] + fn with_distributed_optimize_shuffle_partitioning_ratio(mut self, ratio: f64) -> Result; + fn set_distributed_compression(&mut self, compression: Option) -> Result<(), DataFusionError>; #[call(set_distributed_compression)] #[expr($?;Ok(self))] @@ -998,6 +1038,11 @@ impl DistributedExt for SessionContext { #[expr($?;Ok(self))] fn with_distributed_optimize_shuffle_partitioning(self, enabled: bool) -> Result; + fn set_distributed_optimize_shuffle_partitioning_ratio(&mut self, ratio: f64) -> Result<(), DataFusionError>; + #[call(set_distributed_optimize_shuffle_partitioning_ratio)] + #[expr($?;Ok(self))] + fn with_distributed_optimize_shuffle_partitioning_ratio(self, ratio: f64) -> Result; + fn set_distributed_compression(&mut self, compression: Option) -> Result<(), DataFusionError>; #[call(set_distributed_compression)] #[expr($?;Ok(self))] diff --git a/src/distributed_planner/distributed_config.rs b/src/distributed_planner/distributed_config.rs index 7698b695..fbbbf9b2 100644 --- a/src/distributed_planner/distributed_config.rs +++ b/src/distributed_planner/distributed_config.rs @@ -47,6 +47,11 @@ extensions_options! { pub broadcast_joins: bool, default = false /// Use 1:1 task-partition mapping instead of N*M fanout, eliminating partition inflation. pub optimize_shuffle_partitioning: bool, default = false + /// Scale factor controlling when optimized shuffle partitioning is applied at a boundary. + /// The optimization is applied when cardinality_task_count <= partition_count * ratio. + /// A ratio of 1.0 (default) only optimizes when task count <= partition count exactly. + /// A ratio > 1.0 also optimizes when task count is slightly above partition count. + pub optimize_shuffle_partitioning_ratio: f64, default = 1.0 /// The compression used for sending data over the network between workers. /// It can be set to either `zstd`, `lz4` or `none`. pub compression: String, default = "lz4".to_string() diff --git a/src/distributed_planner/plan_annotator.rs b/src/distributed_planner/plan_annotator.rs index 7b5338d7..7d84cae4 100644 --- a/src/distributed_planner/plan_annotator.rs +++ b/src/distributed_planner/plan_annotator.rs @@ -402,7 +402,8 @@ fn _annotate_plan( annotation.children.first().map(|c| &c.plan_or_nb) { let partition_count = repartition.output_partitioning().partition_count(); - if annotation.task_count.as_usize() <= partition_count { + let threshold = (partition_count as f64 * d_cfg.optimize_shuffle_partitioning_ratio).ceil() as usize; + if annotation.task_count.as_usize() <= threshold { annotation.task_count = Desired(partition_count); } } From 0d2a2e36b34e35b63551e556f1aa2afdaa8fc435 Mon Sep 17 00:00:00 2001 From: Lavkesh Lahngir Date: Fri, 17 Apr 2026 02:39:56 +0200 Subject: [PATCH 7/8] add threshold for optimizer --- src/distributed_planner/plan_annotator.rs | 27 +++++++++++++++++------ 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/src/distributed_planner/plan_annotator.rs b/src/distributed_planner/plan_annotator.rs index 7d84cae4..88e4c0f6 100644 --- a/src/distributed_planner/plan_annotator.rs +++ b/src/distributed_planner/plan_annotator.rs @@ -393,19 +393,32 @@ fn _annotate_plan( // partition count. Override the cardinality-derived downstream task count here to enforce // that invariant so stage scheduling spawns the right number of tasks. // - // Only apply when the cardinality-derived task count does not exceed the upstream partition - // count. If cardinality says more tasks are needed (high-volume stages), fall back to - // non-optimized so parallelism is preserved. + // The optimize_shuffle_partitioning_ratio scales the task count at shuffle boundaries. + // capped = ceil(partition_count * ratio), minimum 1. + // - ratio >= 1.0: clamp(cardinality, partition_count, capped) + // - cardinality < p → lifted to p (optimized 1:1 mapping) + // - p <= cardinality <= capped → kept as-is + // - cardinality > capped → reduced to capped (between p and p*N) + // - ratio=1 always produces p tasks (fully optimized) + // - ratio < 1.0: cardinality.min(capped) — scales down below p, no floor if d_cfg.optimize_shuffle_partitioning { if let PlanOrNetworkBoundary::Shuffle = &annotation.plan_or_nb { if let Some(PlanOrNetworkBoundary::Plan(repartition)) = annotation.children.first().map(|c| &c.plan_or_nb) { + let ratio = d_cfg.optimize_shuffle_partitioning_ratio; let partition_count = repartition.output_partitioning().partition_count(); - let threshold = (partition_count as f64 * d_cfg.optimize_shuffle_partitioning_ratio).ceil() as usize; - if annotation.task_count.as_usize() <= threshold { - annotation.task_count = Desired(partition_count); - } + let capped = ((partition_count as f64 * ratio).ceil() as usize).max(1); + let cardinality = annotation.task_count.as_usize(); + let new_task_count = if ratio >= 1.0 { + // Clamp always applies: lifts below p, caps above capped. + // ratio=1 → task_count=p (optimized). ratio>1 → between p and p*ratio. + cardinality.max(partition_count).min(capped) + } else { + // Scale down below p. No floor — ratio=0.5 → max p/2 tasks. + cardinality.min(capped) + }; + annotation.task_count = Desired(new_task_count); } } } From 4843286c293c6947396496925e68d85739acaf35 Mon Sep 17 00:00:00 2001 From: Lavkesh Lahngir Date: Thu, 23 Apr 2026 14:39:35 +0200 Subject: [PATCH 8/8] fix --- src/execution_plans/network_shuffle.rs | 4 ++-- src/protobuf/distributed_codec.rs | 1 + src/test_utils/plans.rs | 3 ++- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/execution_plans/network_shuffle.rs b/src/execution_plans/network_shuffle.rs index a2b15782..466fd6b8 100644 --- a/src/execution_plans/network_shuffle.rs +++ b/src/execution_plans/network_shuffle.rs @@ -164,9 +164,9 @@ impl NetworkShuffleExec { })?; let properties = if optimize_shuffle_partitioning { - let mut props = input.properties().clone(); + let mut props = input.properties().as_ref().clone(); props.partitioning = Partitioning::UnknownPartitioning(1); - props + Arc::new(props) } else { input.properties().clone() }; diff --git a/src/protobuf/distributed_codec.rs b/src/protobuf/distributed_codec.rs index c25f8977..fc38fff5 100644 --- a/src/protobuf/distributed_codec.rs +++ b/src/protobuf/distributed_codec.rs @@ -723,6 +723,7 @@ mod tests { part, schema, dummy_stage_with_plan(), + false, )); let mut buf = Vec::new(); diff --git a/src/test_utils/plans.rs b/src/test_utils/plans.rs index 9e372fd9..0393506f 100644 --- a/src/test_utils/plans.rs +++ b/src/test_utils/plans.rs @@ -5,8 +5,9 @@ use crate::execution_plans::DistributedExec; use crate::stage::Stage; use crate::test_utils::in_memory_channel_resolver::InMemoryWorkerResolver; use crate::worker::generated::worker::TaskKey; +use crate::DistributedConfig; #[cfg(test)] -use crate::{DistributedConfig, TaskEstimation, TaskEstimator}; +use crate::{TaskEstimation, TaskEstimator}; #[cfg(test)] use datafusion::config::ConfigOptions; use datafusion::{