Revert "Rollback plan dependent usize (#201)"

This reverts commit 01b0866.

Author: jayshrivastava

src/distributed_ext.rs

@@ -1,4 +1,3 @@
-use crate::ChannelResolver;
 use crate::channel_resolver_ext::set_distributed_channel_resolver;
 use crate::config_extension_ext::{
     set_distributed_option_extension, set_distributed_option_extension_from_headers,
@@ -7,6 +6,7 @@ use crate::distributed_planner::{
     set_distributed_network_coalesce_tasks, set_distributed_network_shuffle_tasks,
 };
 use crate::protobuf::{set_distributed_user_codec, set_distributed_user_codec_arc};
+use crate::{ChannelResolver, IntoPlanDependentUsize};
 use datafusion::common::DataFusionError;
 use datafusion::config::ConfigExtension;
 use datafusion::execution::{SessionState, SessionStateBuilder};
@@ -233,10 +233,10 @@ pub trait DistributedExt: Sized {
     ///  ( task 1 )  ( task 2 ) ( task 3 )  N tasks
     /// ```
     /// This parameter defines N
-    fn with_distributed_network_coalesce_tasks(self, tasks: usize) -> Self;
+    fn with_distributed_network_coalesce_tasks<T: IntoPlanDependentUsize>(self, tasks: T) -> Self;
 
     /// Same as [DistributedExt::with_distributed_network_coalesce_tasks] but with an in-place mutation.
-    fn set_distributed_network_coalesce_tasks(&mut self, tasks: usize);
+    fn set_distributed_network_coalesce_tasks<T: IntoPlanDependentUsize>(&mut self, tasks: T);
 
     /// Upon shuffling data, this defines how many tasks are employed into performing the shuffling.
     /// ```text
@@ -246,10 +246,10 @@ pub trait DistributedExt: Sized {
     ///       ( task 1 )  ( task 2 )       N tasks
     /// ```
     /// This parameter defines N
-    fn with_distributed_network_shuffle_tasks(self, tasks: usize) -> Self;
+    fn with_distributed_network_shuffle_tasks<T: IntoPlanDependentUsize>(self, tasks: T) -> Self;
 
     /// Same as [DistributedExt::with_distributed_network_shuffle_tasks] but with an in-place mutation.
-    fn set_distributed_network_shuffle_tasks(&mut self, tasks: usize);
+    fn set_distributed_network_shuffle_tasks<T: IntoPlanDependentUsize>(&mut self, tasks: T);
 }
 
 impl DistributedExt for SessionConfig {
@@ -282,11 +282,11 @@ impl DistributedExt for SessionConfig {
         set_distributed_channel_resolver(self, resolver)
     }
 
-    fn set_distributed_network_coalesce_tasks(&mut self, tasks: usize) {
+    fn set_distributed_network_coalesce_tasks<T: IntoPlanDependentUsize>(&mut self, tasks: T) {
         set_distributed_network_coalesce_tasks(self, tasks)
     }
 
-    fn set_distributed_network_shuffle_tasks(&mut self, tasks: usize) {
+    fn set_distributed_network_shuffle_tasks<T: IntoPlanDependentUsize>(&mut self, tasks: T) {
         set_distributed_network_shuffle_tasks(self, tasks)
     }
 
@@ -314,11 +314,11 @@ impl DistributedExt for SessionConfig {
 
             #[call(set_distributed_network_coalesce_tasks)]
             #[expr($;self)]
-            fn with_distributed_network_coalesce_tasks(mut self, tasks: usize) -> Self;
+            fn with_distributed_network_coalesce_tasks<T: IntoPlanDependentUsize>(mut self, tasks: T) -> Self;
 
             #[call(set_distributed_network_shuffle_tasks)]
             #[expr($;self)]
-            fn with_distributed_network_shuffle_tasks(mut self, tasks: usize) -> Self;
+            fn with_distributed_network_shuffle_tasks<T: IntoPlanDependentUsize>(mut self, tasks: T) -> Self;
         }
     }
 }
@@ -351,15 +351,15 @@ impl DistributedExt for SessionStateBuilder {
             #[expr($;self)]
             fn with_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
 
-            fn set_distributed_network_coalesce_tasks(&mut self, tasks: usize);
+            fn set_distributed_network_coalesce_tasks<T: IntoPlanDependentUsize>(&mut self, tasks: T);
             #[call(set_distributed_network_coalesce_tasks)]
             #[expr($;self)]
-            fn with_distributed_network_coalesce_tasks(mut self, tasks: usize) -> Self;
+            fn with_distributed_network_coalesce_tasks<T: IntoPlanDependentUsize>(mut self, tasks: T) -> Self;
 
-            fn set_distributed_network_shuffle_tasks(&mut self, tasks: usize);
+            fn set_distributed_network_shuffle_tasks<T: IntoPlanDependentUsize>(&mut self, tasks: T);
             #[call(set_distributed_network_shuffle_tasks)]
             #[expr($;self)]
-            fn with_distributed_network_shuffle_tasks(mut self, tasks: usize) -> Self;
+            fn with_distributed_network_shuffle_tasks<T: IntoPlanDependentUsize>(mut self, tasks: T) -> Self;
         }
     }
 }
@@ -392,15 +392,15 @@ impl DistributedExt for SessionState {
             #[expr($;self)]
             fn with_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
 
-            fn set_distributed_network_coalesce_tasks(&mut self, tasks: usize);
+            fn set_distributed_network_coalesce_tasks<T: IntoPlanDependentUsize>(&mut self, tasks: T);
             #[call(set_distributed_network_coalesce_tasks)]
             #[expr($;self)]
-            fn with_distributed_network_coalesce_tasks(mut self, tasks: usize) -> Self;
+            fn with_distributed_network_coalesce_tasks<T: IntoPlanDependentUsize>(mut self, tasks: T) -> Self;
 
-            fn set_distributed_network_shuffle_tasks(&mut self, tasks: usize);
+            fn set_distributed_network_shuffle_tasks<T: IntoPlanDependentUsize>(&mut self, tasks: T);
             #[call(set_distributed_network_shuffle_tasks)]
             #[expr($;self)]
-            fn with_distributed_network_shuffle_tasks(mut self, tasks: usize) -> Self;
+            fn with_distributed_network_shuffle_tasks<T: IntoPlanDependentUsize>(mut self, tasks: T) -> Self;
         }
     }
 }
@@ -433,15 +433,15 @@ impl DistributedExt for SessionContext {
             #[expr($;self)]
             fn with_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(self, resolver: T) -> Self;
 
-            fn set_distributed_network_coalesce_tasks(&mut self, tasks: usize);
+            fn set_distributed_network_coalesce_tasks<T: IntoPlanDependentUsize>(&mut self, tasks: T);
             #[call(set_distributed_network_coalesce_tasks)]
             #[expr($;self)]
-            fn with_distributed_network_coalesce_tasks(self, tasks: usize) -> Self;
+            fn with_distributed_network_coalesce_tasks<T: IntoPlanDependentUsize>(self, tasks: T) -> Self;
 
-            fn set_distributed_network_shuffle_tasks(&mut self, tasks: usize);
+            fn set_distributed_network_shuffle_tasks<T: IntoPlanDependentUsize>(&mut self, tasks: T);
             #[call(set_distributed_network_shuffle_tasks)]
             #[expr($;self)]
-            fn with_distributed_network_shuffle_tasks(self, tasks: usize) -> Self;
+            fn with_distributed_network_shuffle_tasks<T: IntoPlanDependentUsize>(self, tasks: T) -> Self;
         }
     }
 }

src/distributed_planner/distributed_config.rs

@@ -1,7 +1,66 @@
 use datafusion::common::extensions_options;
-use datafusion::config::ConfigExtension;
+use datafusion::config::{ConfigExtension, ConfigField, Visit, default_config_transform};
+use datafusion::physical_plan::ExecutionPlan;
 use datafusion::prelude::SessionConfig;
-use std::fmt::Debug;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::Arc;
+
+#[derive(Clone)]
+#[allow(clippy::type_complexity)]
+pub struct PlanDependentUsize(
+    pub(crate) Arc<dyn Fn(&Arc<dyn ExecutionPlan>) -> usize + Send + Sync>,
+);
+
+impl PlanDependentUsize {
+    pub fn call(&self, plan: &Arc<dyn ExecutionPlan>) -> usize {
+        self.0(plan)
+    }
+}
+
+pub trait IntoPlanDependentUsize {
+    fn into_plan_dependent_usize(self) -> PlanDependentUsize;
+}
+
+impl IntoPlanDependentUsize for usize {
+    fn into_plan_dependent_usize(self) -> PlanDependentUsize {
+        PlanDependentUsize(Arc::new(move |_| self))
+    }
+}
+
+impl<T: Fn(&Arc<dyn ExecutionPlan>) -> usize + Send + Sync + 'static> IntoPlanDependentUsize for T {
+    fn into_plan_dependent_usize(self) -> PlanDependentUsize {
+        PlanDependentUsize(Arc::new(self))
+    }
+}
+
+impl Default for PlanDependentUsize {
+    fn default() -> Self {
+        PlanDependentUsize(Arc::new(|_| 0))
+    }
+}
+
+impl Debug for PlanDependentUsize {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "PlanDependantUsize")
+    }
+}
+
+impl Display for PlanDependentUsize {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "PlanDependantUsize")
+    }
+}
+
+impl ConfigField for PlanDependentUsize {
+    fn visit<V: Visit>(&self, v: &mut V, key: &str, description: &'static str) {
+        v.some(key, self, description);
+    }
+
+    fn set(&mut self, _: &str, value: &str) -> datafusion::common::Result<()> {
+        *self = default_config_transform::<usize>(value)?.into_plan_dependent_usize();
+        Ok(())
+    }
+}
 
 extensions_options! {
     pub struct DistributedConfig {
@@ -13,7 +72,7 @@ extensions_options! {
         ///       ( task 1 )  ( task 2 )       N tasks
         /// ```
         /// This parameter defines N
-        pub network_shuffle_tasks: Option<usize>, default = None
+        pub network_shuffle_tasks: Option<PlanDependentUsize>, default = None
         /// Upon merging multiple tasks into one, this defines how many tasks are merged.
         /// ```text
         ///              ( task 1 )
@@ -22,7 +81,7 @@ extensions_options! {
         ///  ( task 1 )  ( task 2 ) ( task 3 )  N tasks
         /// ```
         /// This parameter defines N
-        pub network_coalesce_tasks: Option<usize>, default = None
+        pub network_coalesce_tasks: Option<PlanDependentUsize>, default = None
     }
 }
 
@@ -32,30 +91,36 @@ impl ConfigExtension for DistributedConfig {
 
 impl DistributedConfig {
     /// Sets the amount of tasks used in a network shuffle operation.
-    pub fn with_network_shuffle_tasks(mut self, tasks: usize) -> Self {
-        self.network_shuffle_tasks = Some(tasks);
+    pub fn with_network_shuffle_tasks(mut self, tasks: impl IntoPlanDependentUsize) -> Self {
+        self.network_shuffle_tasks = Some(tasks.into_plan_dependent_usize());
         self
     }
 
     /// Sets the amount of tasks used in a network coalesce operation.
-    pub fn with_network_coalesce_tasks(mut self, tasks: usize) -> Self {
-        self.network_coalesce_tasks = Some(tasks);
+    pub fn with_network_coalesce_tasks(mut self, tasks: impl IntoPlanDependentUsize) -> Self {
+        self.network_coalesce_tasks = Some(tasks.into_plan_dependent_usize());
         self
     }
 }
 
-pub(crate) fn set_distributed_network_coalesce_tasks(cfg: &mut SessionConfig, tasks: usize) {
+pub(crate) fn set_distributed_network_coalesce_tasks(
+    cfg: &mut SessionConfig,
+    tasks: impl IntoPlanDependentUsize,
+) {
     let ext = &mut cfg.options_mut().extensions;
     let Some(prev) = ext.get_mut::<DistributedConfig>() else {
         return ext.insert(DistributedConfig::default().with_network_coalesce_tasks(tasks));
     };
-    prev.network_coalesce_tasks = Some(tasks);
+    prev.network_coalesce_tasks = Some(tasks.into_plan_dependent_usize());
 }
 
-pub(crate) fn set_distributed_network_shuffle_tasks(cfg: &mut SessionConfig, tasks: usize) {
+pub(crate) fn set_distributed_network_shuffle_tasks(
+    cfg: &mut SessionConfig,
+    tasks: impl IntoPlanDependentUsize,
+) {
     let ext = &mut cfg.options_mut().extensions;
     let Some(prev) = ext.get_mut::<DistributedConfig>() else {
         return ext.insert(DistributedConfig::default().with_network_shuffle_tasks(tasks));
     };
-    prev.network_shuffle_tasks = Some(tasks);
+    prev.network_shuffle_tasks = Some(tasks.into_plan_dependent_usize());
 }

src/distributed_planner/distributed_physical_optimizer_rule.rs

@@ -110,13 +110,14 @@ pub fn apply_network_boundaries(
         }
 
         // If this is a hash RepartitionExec, introduce a shuffle.
-        if let (Some(node), Some(input_tasks)) = (
+        if let (Some(node), Some(tasks)) = (
             plan.as_any().downcast_ref::<RepartitionExec>(),
-            cfg.network_shuffle_tasks,
+            cfg.network_shuffle_tasks.clone(),
         ) {
             if !matches!(node.partitioning(), Partitioning::Hash(_, _)) {
                 return Ok(Transformed::no(plan));
             }
+            let input_tasks = tasks.0(&plan);
             if input_tasks == 0 {
                 return Ok(Transformed::no(plan));
             }
@@ -128,9 +129,9 @@ pub fn apply_network_boundaries(
         // If this is a CoalescePartitionsExec, it means that the original plan is trying to
         // merge all partitions into one. We need to go one step ahead and also merge all tasks
         // into one.
-        if let (Some(node), Some(input_tasks)) = (
+        if let (Some(node), Some(tasks)) = (
             plan.as_any().downcast_ref::<CoalescePartitionsExec>(),
-            cfg.network_coalesce_tasks,
+            cfg.network_coalesce_tasks.clone(),
         ) {
             // If the immediate child is a PartitionIsolatorExec, it means that the rest of the
             // plan is just a couple of non-computational nodes that are probably not worth
@@ -139,6 +140,7 @@ pub fn apply_network_boundaries(
                 return Ok(Transformed::no(plan));
             }
 
+            let input_tasks = tasks.0(&plan);
             if input_tasks == 0 {
                 return Ok(Transformed::no(plan));
             }
@@ -152,10 +154,11 @@ pub fn apply_network_boundaries(
         // The SortPreservingMergeExec node will try to coalesce all partitions into just 1.
         // We need to account for it and help it by also coalescing all tasks into one, therefore
         // a NetworkCoalesceExec is introduced.
-        if let (Some(node), Some(input_tasks)) = (
+        if let (Some(node), Some(tasks)) = (
             plan.as_any().downcast_ref::<SortPreservingMergeExec>(),
-            cfg.network_coalesce_tasks,
+            cfg.network_coalesce_tasks.clone(),
         ) {
+            let input_tasks = tasks.0(&plan);
             if input_tasks == 0 {
                 return Ok(Transformed::no(plan));
             }

src/distributed_planner/mod.rs

@@ -7,7 +7,7 @@ pub(crate) use distributed_config::{
     set_distributed_network_coalesce_tasks, set_distributed_network_shuffle_tasks,
 };
 
-pub use distributed_config::DistributedConfig;
+pub use distributed_config::{DistributedConfig, IntoPlanDependentUsize, PlanDependentUsize};
 pub use distributed_physical_optimizer_rule::{
     DistributedPhysicalOptimizerRule, apply_network_boundaries, distribute_plan,
 };

src/lib.rs

@@ -17,8 +17,9 @@ pub mod test_utils;
 pub use channel_resolver_ext::{BoxCloneSyncChannel, ChannelResolver};
 pub use distributed_ext::DistributedExt;
 pub use distributed_planner::{
-    DistributedConfig, DistributedPhysicalOptimizerRule, InputStageInfo, NetworkBoundary,
-    NetworkBoundaryExt, apply_network_boundaries, distribute_plan,
+    DistributedConfig, DistributedPhysicalOptimizerRule, InputStageInfo, IntoPlanDependentUsize,
+    NetworkBoundary, NetworkBoundaryExt, PlanDependentUsize, apply_network_boundaries,
+    distribute_plan,
 };
 pub use execution_plans::{
     DistributedExec, NetworkCoalesceExec, NetworkShuffleExec, PartitionIsolatorExec,