Refactor to NetworkShuffleExec and NetworkCoalesceExec

Author: gabotechs

benchmarks/src/tpch/run.rs

@@ -143,8 +143,8 @@ impl DistributedSessionBuilder for RunOpt {
         if !self.workers.is_empty() {
             let tasks = self.max_tasks.unwrap_or(self.workers.len());
             let rule = DistributedPhysicalOptimizerRule::new()
-                .with_coalesce_partitions_exec_tasks(tasks)
-                .with_network_shuffle_exec_tasks(tasks);
+                .with_network_coalesce_tasks(tasks)
+                .with_network_shuffle_tasks(tasks);
             builder = builder.with_physical_optimizer_rule(Arc::new(rule));
         }
 

src/common/mod.rs

@@ -1,5 +1,7 @@
 mod composed_extension_codec;
+mod partitioning;
 #[allow(unused)]
 pub mod ttl_map;
 
 pub(crate) use composed_extension_codec::ComposedPhysicalExtensionCodec;
+pub(crate) use partitioning::{scale_partitioning, scale_partitioning_props};

src/common/partitioning.rs

@@ -0,0 +1,25 @@
+use datafusion::physical_expr::Partitioning;
+use datafusion::physical_plan::PlanProperties;
+
+pub fn scale_partitioning_props(
+    props: &PlanProperties,
+    f: impl FnOnce(usize) -> usize,
+) -> PlanProperties {
+    PlanProperties::new(
+        props.eq_properties.clone(),
+        scale_partitioning(&props.partitioning, f),
+        props.emission_type,
+        props.boundedness,
+    )
+}
+
+pub fn scale_partitioning(
+    partitioning: &Partitioning,
+    f: impl FnOnce(usize) -> usize,
+) -> Partitioning {
+    match &partitioning {
+        Partitioning::RoundRobinBatch(p) => Partitioning::RoundRobinBatch(f(*p)),
+        Partitioning::Hash(hash, p) => Partitioning::Hash(hash.clone(), f(*p)),
+        Partitioning::UnknownPartitioning(p) => Partitioning::UnknownPartitioning(f(*p)),
+    }
+}

src/distributed_physical_optimizer_rule.rs

@@ -1,9 +1,9 @@
-mod network_coalesce_tasks;
-mod network_hash_shuffle;
+mod network_coalesce;
+mod network_shuffle;
 mod partition_isolator;
 mod stage;
 
-pub use network_coalesce_tasks::{NetworkCoalesceTasksExec, NetworkCoalesceTasksReadyExec};
-pub use network_hash_shuffle::{NetworkHashShuffleExec, NetworkHashShuffleReadyExec};
+pub use network_coalesce::{NetworkCoalesceExec, NetworkCoalesceReady};
+pub use network_shuffle::{NetworkShuffleExec, NetworkShuffleReadyExec};
 pub use partition_isolator::PartitionIsolatorExec;
 pub use stage::{DistributedTaskContext, ExecutionTask, StageExec};

src/execution_plans/mod.rs

@@ -1,4 +1,5 @@
 use crate::channel_resolver_ext::get_distributed_channel_resolver;
+use crate::common::scale_partitioning_props;
 use crate::config_extension_ext::ContextGrpcMetadata;
 use crate::distributed_physical_optimizer_rule::{limit_tasks_err, DistributedExecutionPlan};
 use crate::errors::{map_flight_to_datafusion_error, map_status_to_datafusion_error};
@@ -13,7 +14,6 @@ use arrow_flight::Ticket;
 use datafusion::common::{exec_err, internal_datafusion_err, internal_err, plan_err};
 use datafusion::error::DataFusionError;
 use datafusion::execution::{SendableRecordBatchStream, TaskContext};
-use datafusion::physical_expr::Partitioning;
 use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
@@ -32,34 +32,34 @@ use tonic::Request;
 /// 2. Ready: runs within a distributed stage and queries the next input stage over the network
 ///     using Arrow Flight.
 #[derive(Debug, Clone)]
-pub enum NetworkCoalesceTasksExec {
-    Pending(NetworkCoalesceTasksPendingExec),
-    Ready(NetworkCoalesceTasksReadyExec),
+pub enum NetworkCoalesceExec {
+    Pending(NetworkCoalescePending),
+    Ready(NetworkCoalesceReady),
 }
 
-/// Placeholder version of the [NetworkCoalesceTasksExec] node. It acts as a marker for the
+/// Placeholder version of the [NetworkCoalesceExec] node. It acts as a marker for the
 /// distributed optimization step, which will replace it with the appropriate
-/// [NetworkCoalesceTasksReadyExec] node.
+/// [NetworkCoalesceReady] node.
 #[derive(Debug, Clone)]
-pub struct NetworkCoalesceTasksPendingExec {
+pub struct NetworkCoalescePending {
     properties: PlanProperties,
     input_tasks: usize,
     child: Arc<dyn ExecutionPlan>,
 }
 
-/// Ready version of the [NetworkCoalesceTasksExec] node. This node can be created in
+/// Ready version of the [NetworkCoalesceExec] node. This node can be created in
 /// just two ways:
-/// - by the distributed optimization step based on an original [NetworkCoalesceTasksPendingExec]
+/// - by the distributed optimization step based on an original [NetworkCoalescePending]
 /// - deserialized from a protobuf plan sent over the network.
 #[derive(Debug, Clone)]
-pub struct NetworkCoalesceTasksReadyExec {
+pub struct NetworkCoalesceReady {
     /// the properties we advertise for this execution plan
     pub(crate) properties: PlanProperties,
     pub(crate) stage_num: usize,
     pub(crate) input_tasks: usize,
 }
 
-impl NetworkCoalesceTasksExec {
+impl NetworkCoalesceExec {
     pub fn from_coalesce_partitions_exec(
         input: &CoalescePartitionsExec,
         input_tasks: usize,
@@ -83,15 +83,15 @@ impl NetworkCoalesceTasksExec {
             return internal_err!("Expected a single child");
         };
 
-        Ok(Self::Pending(NetworkCoalesceTasksPendingExec {
+        Ok(Self::Pending(NetworkCoalescePending {
             properties: child.properties().clone(),
             input_tasks,
             child: Arc::clone(child),
         }))
     }
 }
 
-impl DistributedExecutionPlan for NetworkCoalesceTasksExec {
+impl DistributedExecutionPlan for NetworkCoalesceExec {
     fn to_stage_info(
         &self,
         n_tasks: usize,
@@ -112,12 +112,14 @@ impl DistributedExecutionPlan for NetworkCoalesceTasksExec {
         stage_num: usize,
         stage_head: &Arc<dyn ExecutionPlan>,
     ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
-        let NetworkCoalesceTasksExec::Pending(pending) = self else {
-            return internal_err!("NetworkCoalesceTasksExec is already distributed");
+        let NetworkCoalesceExec::Pending(pending) = self else {
+            return internal_err!("NetworkCoalesceExec is already distributed");
         };
 
-        let ready = NetworkCoalesceTasksReadyExec {
-            properties: scale_partitioning(stage_head.properties(), |p| p * pending.input_tasks),
+        let ready = NetworkCoalesceReady {
+            properties: scale_partitioning_props(stage_head.properties(), |p| {
+                p * pending.input_tasks
+            }),
             stage_num,
             input_tasks: pending.input_tasks,
         };
@@ -127,49 +129,33 @@ impl DistributedExecutionPlan for NetworkCoalesceTasksExec {
 
     fn with_input_tasks(&self, input_tasks: usize) -> Arc<dyn DistributedExecutionPlan> {
         Arc::new(match self {
-            NetworkCoalesceTasksExec::Pending(pending) => {
-                NetworkCoalesceTasksExec::Pending(NetworkCoalesceTasksPendingExec {
+            NetworkCoalesceExec::Pending(pending) => {
+                NetworkCoalesceExec::Pending(NetworkCoalescePending {
                     properties: pending.properties.clone(),
                     input_tasks,
                     child: pending.child.clone(),
                 })
             }
-            NetworkCoalesceTasksExec::Ready(ready) => {
-                NetworkCoalesceTasksExec::Ready(NetworkCoalesceTasksReadyExec {
-                    properties: scale_partitioning(&ready.properties, |p| {
-                        p * input_tasks / ready.input_tasks
-                    }),
-                    stage_num: ready.stage_num,
-                    input_tasks,
-                })
-            }
+            NetworkCoalesceExec::Ready(ready) => NetworkCoalesceExec::Ready(NetworkCoalesceReady {
+                properties: scale_partitioning_props(&ready.properties, |p| {
+                    p * input_tasks / ready.input_tasks
+                }),
+                stage_num: ready.stage_num,
+                input_tasks,
+            }),
         })
     }
 }
 
-fn scale_partitioning(props: &PlanProperties, f: impl FnOnce(usize) -> usize) -> PlanProperties {
-    let partitioning = match &props.partitioning {
-        Partitioning::RoundRobinBatch(p) => Partitioning::RoundRobinBatch(f(*p)),
-        Partitioning::Hash(hash, p) => Partitioning::Hash(hash.clone(), f(*p)),
-        Partitioning::UnknownPartitioning(p) => Partitioning::UnknownPartitioning(f(*p)),
-    };
-    PlanProperties::new(
-        props.eq_properties.clone(),
-        partitioning,
-        props.emission_type,
-        props.boundedness,
-    )
-}
-
-impl DisplayAs for NetworkCoalesceTasksExec {
+impl DisplayAs for NetworkCoalesceExec {
     fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
-        write!(f, "NetworkCoalesceTasksExec")
+        write!(f, "NetworkCoalesceExec")
     }
 }
 
-impl ExecutionPlan for NetworkCoalesceTasksExec {
+impl ExecutionPlan for NetworkCoalesceExec {
     fn name(&self) -> &str {
-        "NetworkCoalesceTasksExec"
+        "NetworkCoalesceExec"
     }
 
     fn as_any(&self) -> &dyn Any {
@@ -178,15 +164,15 @@ impl ExecutionPlan for NetworkCoalesceTasksExec {
 
     fn properties(&self) -> &PlanProperties {
         match self {
-            NetworkCoalesceTasksExec::Pending(v) => &v.properties,
-            NetworkCoalesceTasksExec::Ready(v) => &v.properties,
+            NetworkCoalesceExec::Pending(v) => &v.properties,
+            NetworkCoalesceExec::Ready(v) => &v.properties,
         }
     }
 
     fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
         match self {
-            NetworkCoalesceTasksExec::Pending(v) => vec![&v.child],
-            NetworkCoalesceTasksExec::Ready(_) => vec![],
+            NetworkCoalesceExec::Pending(v) => vec![&v.child],
+            NetworkCoalesceExec::Ready(_) => vec![],
         }
     }
 
@@ -196,7 +182,7 @@ impl ExecutionPlan for NetworkCoalesceTasksExec {
     ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
         if !children.is_empty() {
             return plan_err!(
-                "NetworkCoalesceTasksExec: wrong number of children, expected 0, got {}",
+                "NetworkCoalesceExec: wrong number of children, expected 0, got {}",
                 children.len()
             );
         }
@@ -208,16 +194,16 @@ impl ExecutionPlan for NetworkCoalesceTasksExec {
         partition: usize,
         context: Arc<TaskContext>,
     ) -> Result<SendableRecordBatchStream, DataFusionError> {
-        let NetworkCoalesceTasksExec::Ready(self_ready) = self else {
+        let NetworkCoalesceExec::Ready(self_ready) = self else {
             return exec_err!(
-                "NetworkCoalesceTasksExec is not ready, was the distributed optimization step performed?"
+                "NetworkCoalesceExec is not ready, was the distributed optimization step performed?"
             );
         };
 
         // get the channel manager and current stage from our context
         let channel_resolver = get_distributed_channel_resolver(context.session_config())?;
 
-        // the `NetworkCoalesceTasksExec` node can only be executed in the context of a `StageExec`
+        // the `NetworkCoalesceExec` node can only be executed in the context of a `StageExec`
         let stage = StageExec::from_ctx(&context)?;
 
         // of our child stages find the one that matches the one we are supposed to be
@@ -226,15 +212,13 @@ impl ExecutionPlan for NetworkCoalesceTasksExec {
 
         let codec = DistributedCodec::new_combined_with_user(context.session_config());
         let child_stage_proto = proto_from_stage(child_stage, &codec).map_err(|e| {
-            internal_datafusion_err!(
-                "NetworkCoalesceTasksExec: failed to convert stage to proto: {e}"
-            )
+            internal_datafusion_err!("NetworkCoalesceExec: failed to convert stage to proto: {e}")
         })?;
 
         let context_headers = ContextGrpcMetadata::headers_from_ctx(&context);
         let task_context = DistributedTaskContext::from_ctx(&context);
         if task_context.task_index > 0 {
-            return exec_err!("NetworkCoalesceTasksExec cannot be executed in more than one task");
+            return exec_err!("NetworkCoalesceExec cannot be executed in more than one task");
         }
 
         let partitions_per_task =
@@ -267,7 +251,7 @@ impl ExecutionPlan for NetworkCoalesceTasksExec {
         };
 
         let Some(url) = task.url.clone() else {
-            return internal_err!("NetworkCoalesceTasksExec: task is unassigned, cannot proceed");
+            return internal_err!("NetworkCoalesceExec: task is unassigned, cannot proceed");
         };
 
         let stream = async move {