Lavkesh/network shuffle experiments

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<Self, DataFusionError>;
+
+    /// 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<Self, DataFusionError>;
+
+    /// 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, DataFusionError> {
+        self.set_distributed_optimize_shuffle_partitioning_ratio(ratio)?;
+        Ok(self)
+    }
+
     fn set_distributed_compression(
         &mut self,
         compression: Option<CompressionType>,
@@ -703,6 +755,10 @@ impl DistributedExt for SessionConfig {
             #[expr($?;Ok(self))]
             fn with_distributed_broadcast_joins(mut self, enabled: bool) -> Result<Self, DataFusionError>;
 
+            #[call(set_distributed_optimize_shuffle_partitioning)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_optimize_shuffle_partitioning(mut self, enabled: bool) -> Result<Self, DataFusionError>;
+
             #[call(set_distributed_compression)]
             #[expr($?;Ok(self))]
             fn with_distributed_compression(mut self, compression: Option<CompressionType>) -> Result<Self, DataFusionError>;
@@ -785,6 +841,16 @@ impl DistributedExt for SessionStateBuilder {
             #[expr($?;Ok(self))]
             fn with_distributed_broadcast_joins(mut self, enabled: bool) -> Result<Self, DataFusionError>;
 
+            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<Self, DataFusionError>;
+
+            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<Self, DataFusionError>;
+
             fn set_distributed_compression(&mut self, compression: Option<CompressionType>) -> 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<Self, DataFusionError>;
 
+            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<Self, DataFusionError>;
+
+            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<Self, DataFusionError>;
+
             fn set_distributed_compression(&mut self, compression: Option<CompressionType>) -> 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<Self, DataFusionError>;
 
+            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<Self, DataFusionError>;
+
+            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<Self, DataFusionError>;
+
             fn set_distributed_compression(&mut self, compression: Option<CompressionType>) -> Result<(), DataFusionError>;
             #[call(set_distributed_compression)]
             #[expr($?;Ok(self))]

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

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()