diff --git a/src/distributed_ext.rs b/src/distributed_ext.rs index 7c957926..8a6c8a1b 100644 --- a/src/distributed_ext.rs +++ b/src/distributed_ext.rs @@ -455,6 +455,32 @@ 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>; + + /// 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]. @@ -615,6 +641,32 @@ 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_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, @@ -703,6 +755,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 +841,16 @@ 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_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))] @@ -871,6 +937,16 @@ 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_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))] @@ -957,6 +1033,16 @@ 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_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/batch_coalescing_below_network_boundaries.rs b/src/distributed_planner/batch_coalescing_below_network_boundaries.rs index 80da921a..238a3525 100644 --- a/src/distributed_planner/batch_coalescing_below_network_boundaries.rs +++ b/src/distributed_planner/batch_coalescing_below_network_boundaries.rs @@ -84,7 +84,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 @@ -112,7 +112,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 @@ -141,7 +141,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_config.rs b/src/distributed_planner/distributed_config.rs index f4dc8284..fbbbf9b2 100644 --- a/src/distributed_planner/distributed_config.rs +++ b/src/distributed_planner/distributed_config.rs @@ -45,6 +45,13 @@ 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 + /// 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/distributed_physical_optimizer_rule.rs b/src/distributed_planner/distributed_physical_optimizer_rule.rs index 40ad3af1..9f83931e 100644 --- a/src/distributed_planner/distributed_physical_optimizer_rule.rs +++ b/src/distributed_planner/distributed_physical_optimizer_rule.rs @@ -118,12 +118,23 @@ 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)?; + + // 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( - require_one_child(new_children)?, + child_plan, query_id, *stage_id, task_count, max_child_task_count.unwrap_or(1), + use_optimized, )?); stage_id.add_assign(1); Ok(node) @@ -237,7 +248,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 @@ -267,7 +278,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 @@ -317,7 +328,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 @@ -370,7 +381,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 @@ -435,12 +446,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 @@ -499,7 +510,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 @@ -963,4 +974,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, optimized=true + └────────────────────────────────────────────────── + ┌───── 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, 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 + │ 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/distributed_planner/plan_annotator.rs b/src/distributed_planner/plan_annotator.rs index 9aeeb461..88e4c0f6 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,42 @@ 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. + // + // 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 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); + } + } + } + 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/benchmarks/shuffle_bench.rs b/src/execution_plans/benchmarks/shuffle_bench.rs index 4e630533..1d74246d 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 14687376..7a45bfe6 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 { return 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, }) } } @@ -188,9 +203,10 @@ 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; + let optimized = self.optimize_shuffle_partitioning; write!( f, - "[Stage {stage}] => NetworkShuffleExec: output_partitions={partitions}, input_tasks={input_tasks}", + "[Stage {stage}] => NetworkShuffleExec: output_partitions={partitions}, input_tasks={input_tasks}, optimized={optimized}", ) } } @@ -230,19 +246,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 +295,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 80fb9eba..1e1c33a5 100644 --- a/src/protobuf/distributed_codec.rs +++ b/src/protobuf/distributed_codec.rs @@ -78,6 +78,7 @@ impl PhysicalExtensionCodec for DistributedCodec { schema, partitioning, input_stage, + optimize_shuffle_partitioning, }) => { let schema: Schema = schema .as_ref() @@ -96,6 +97,7 @@ impl PhysicalExtensionCodec for DistributedCodec { partitioning, Arc::new(schema), parse_stage_proto(input_stage, inputs)?, + optimize_shuffle_partitioning, ))) } DistributedExecNode::NetworkCoalesceTasks(NetworkCoalesceExecProto { @@ -237,6 +239,7 @@ impl PhysicalExtensionCodec for DistributedCodec { &DistributedCodec {}, )?), input_stage: Some(encode_stage_proto(node.input_stage())?), + optimize_shuffle_partitioning: node.optimize_shuffle_partitioning, }; let wrapper = DistributedExecProto { @@ -385,6 +388,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)] @@ -415,6 +420,7 @@ fn new_network_hash_shuffle_exec( partitioning: Partitioning, schema: SchemaRef, input_stage: Stage, + optimize_shuffle_partitioning: bool, ) -> NetworkShuffleExec { NetworkShuffleExec { properties: PlanProperties::new( @@ -426,6 +432,7 @@ fn new_network_hash_shuffle_exec( worker_connections: WorkerConnectionPool::new(input_stage.tasks.len()), input_stage, metrics_collection: Default::default(), + optimize_shuffle_partitioning, } } @@ -565,7 +572,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)?; @@ -586,6 +593,7 @@ mod tests { Partitioning::UnknownPartitioning(1), schema, dummy_stage(), + false, )); let plan: Arc = Arc::new(PartitionIsolatorExec::new(flight.clone(), 1)); @@ -609,11 +617,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()])?; @@ -638,6 +648,7 @@ mod tests { Partitioning::UnknownPartitioning(1), schema.clone(), dummy_stage(), + false, )); let sort_expr = PhysicalSortExpr { @@ -743,11 +754,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 = 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);