use geometric mean 200

Eventual-Inc · Jan 29, 2025 · 062143b · 062143b
1 parent eac7490
commit 062143b
Showing 1 changed file with 19 additions and 9 deletions.
diff --git a/src/daft-physical-plan/src/ops/shuffle_exchange.rs b/src/daft-physical-plan/src/ops/shuffle_exchange.rs
@@ -111,12 +111,22 @@ pub struct ShuffleExchangeFactory {
 }
 
 impl ShuffleExchangeFactory {
-    const PARTITION_THRESHOLD_TO_USE_PRE_SHUFFLE_MERGE: usize = 1_000_000;
+    const PARTITION_THRESHOLD_TO_USE_PRE_SHUFFLE_MERGE: usize = 200;
 
     pub fn new(input: PhysicalPlanRef) -> Self {
         Self { input }
     }
 
+    fn should_use_pre_shuffle_merge(
+        &self,
+        input_num_partitions: usize,
+        target_num_partitions: usize,
+    ) -> bool {
+        let total_num_partitions = input_num_partitions * target_num_partitions;
+        let geometric_mean = (total_num_partitions as f64).sqrt() as usize;
+        geometric_mean > Self::PARTITION_THRESHOLD_TO_USE_PRE_SHUFFLE_MERGE
+    }
+
     fn get_shuffle_strategy(
         &self,
         clustering_spec: Arc<ClusteringSpec>,
@@ -135,10 +145,10 @@ impl ShuffleExchangeFactory {
                 }
             }
             Some(cfg) if cfg.shuffle_algorithm == "auto" => {
-                let intermediate_num_partitions = self.input.clustering_spec().num_partitions()
-                    * clustering_spec.num_partitions();
-                if intermediate_num_partitions > Self::PARTITION_THRESHOLD_TO_USE_PRE_SHUFFLE_MERGE
-                {
+                if self.should_use_pre_shuffle_merge(
+                    self.input.clustering_spec().num_partitions(),
+                    clustering_spec.num_partitions(),
+                ) {
                     ShuffleExchangeStrategy::MapReduceWithPreShuffleMerge {
                         target_spec: clustering_spec,
                         pre_shuffle_merge_threshold: cfg.pre_shuffle_merge_threshold,
@@ -150,10 +160,10 @@ impl ShuffleExchangeFactory {
                 }
             }
             None => {
-                let intermediate_num_partitions = self.input.clustering_spec().num_partitions()
-                    * clustering_spec.num_partitions();
-                if intermediate_num_partitions > Self::PARTITION_THRESHOLD_TO_USE_PRE_SHUFFLE_MERGE
-                {
+                if self.should_use_pre_shuffle_merge(
+                    self.input.clustering_spec().num_partitions(),
+                    clustering_spec.num_partitions(),
+                ) {
                     ShuffleExchangeStrategy::MapReduceWithPreShuffleMerge {
                         target_spec: clustering_spec,
                         pre_shuffle_merge_threshold: DaftExecutionConfig::default()