Rename ArrowFlightReadExec to NetworkShuffleExec

Author: gabotechs

src/execution_plans/metrics.rs

@@ -38,7 +38,7 @@ impl TreeNodeRewriter for TaskMetricsCollector {
     type Node = Arc<dyn ExecutionPlan>;
 
     fn f_down(&mut self, plan: Self::Node) -> Result<Transformed<Self::Node>> {
-        // If the plan is an ArrowFlightReadExec, assume it has collected metrics already
+        // If the plan is an NetworkShuffleExec, assume it has collected metrics already
         // from child tasks.
         let metrics_collection =
             if let Some(node) = plan.as_any().downcast_ref::<NetworkShuffleExec>() {
@@ -79,7 +79,7 @@ impl TreeNodeRewriter for TaskMetricsCollector {
                     }
                 }
             }
-            // Skip the subtree of the ArrowFlightReadExec.
+            // Skip the subtree of the NetworkShuffleExec.
             return Ok(Transformed::new(plan, false, TreeNodeRecursion::Jump));
         }
 
@@ -107,7 +107,7 @@ impl TaskMetricsCollector {
     /// collect metrics from a StageExec plan and any child tasks.
     /// Returns
     /// - a vec representing the metrics for the current task (ordered using a pre-order traversal)
-    /// - a map representing the metrics for some subset of child tasks collected from ArrowFlightReadExec leaves
+    /// - a map representing the metrics for some subset of child tasks collected from NetworkShuffleExec leaves
     #[allow(dead_code)]
     pub fn collect(mut self, stage: &StageExec) -> Result<MetricsCollectorResult, DataFusionError> {
         stage.plan.clone().rewrite(&mut self)?;
@@ -123,14 +123,14 @@ impl TaskMetricsCollector {
 /// Ex. for a plan with the form
 /// AggregateExec
 ///  └── ProjectionExec
-///      └── ArrowFlightReadExec
+///      └── NetworkShuffleExec
 ///
 /// the task will be rewritten as
 ///
 /// MetricsWrapperExec (wrapped: AggregateExec)
 ///  └── MetricsWrapperExec (wrapped: ProjectionExec)
-///      └── ArrowFlightReadExec
-/// (Note that the ArrowFlightReadExec node is not wrapped)
+///      └── NetworkShuffleExec
+/// (Note that the NetworkShuffleExec node is not wrapped)
 pub struct TaskMetricsRewriter {
     metrics: Vec<MetricsSetProto>,
     idx: usize,
@@ -402,7 +402,7 @@ mod tests {
     #[ignore]
     async fn test_metrics_rewriter() {
         let (test_stage, _ctx) = make_test_stage_exec_with_5_nodes().await;
-        let test_metrics_sets = (0..5) // 5 nodes excluding ArrowFlightReadExec
+        let test_metrics_sets = (0..5) // 5 nodes excluding NetworkShuffleExec
             .map(|i| make_distinct_metrics_set(i + 10))
             .collect::<Vec<MetricsSetProto>>();
 
@@ -420,7 +420,7 @@ mod tests {
             r"    ProjectionExec: expr=[id@0 as id, count(Int64(1))@1 as count], metrics=[output_rows=12, elapsed_compute=12ns, start_timestamp=2025-09-18 13:00:12 UTC, end_timestamp=2025-09-18 13:00:13 UTC]",
             r"      AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[count(Int64(1))], metrics=[output_rows=13, elapsed_compute=13ns, start_timestamp=2025-09-18 13:00:13 UTC, end_timestamp=2025-09-18 13:00:14 UTC]",
             r"        CoalesceBatchesExec: target_batch_size=8192, metrics=[output_rows=14, elapsed_compute=14ns, start_timestamp=2025-09-18 13:00:14 UTC, end_timestamp=2025-09-18 13:00:15 UTC]",
-            r"          ArrowFlightReadExec, metrics=[]",
+            r"          NetworkShuffleExec, metrics=[]",
             "" // trailing newline
         ].join("\n");
         assert_eq!(expected, plan_str.to_string());

src/execution_plans/network_coalesce.rs

@@ -95,7 +95,7 @@ pub struct NetworkCoalesceReady {
     ///
     /// An instance may receive metrics for 0 to N child tasks, where N is the number of tasks in
     /// the stage it is reading from. This is because, by convention, the ArrowFlightEndpoint
-    /// sends metrics for a task to the last ArrowFlightReadExec to read from it, which may or may
+    /// sends metrics for a task to the last NetworkCoalesceExec to read from it, which may or may
     /// not be this instance.
     pub(crate) metrics_collection: Arc<DashMap<StageKey, Vec<MetricsSetProto>>>,
 }

src/execution_plans/network_shuffle.rs

@@ -141,7 +141,7 @@ pub struct NetworkShuffleReadyExec {
     ///
     /// An instance may receive metrics for 0 to N child tasks, where N is the number of tasks in
     /// the stage it is reading from. This is because, by convention, the ArrowFlightEndpoint
-    /// sends metrics for a task to the last ArrowFlightReadExec to read from it, which may or may
+    /// sends metrics for a task to the last NetworkShuffleExec to read from it, which may or may
     /// not be this instance.
     pub(crate) metrics_collection: Arc<DashMap<StageKey, Vec<MetricsSetProto>>>,
 }

src/execution_plans/stage.rs

@@ -220,16 +220,14 @@ impl StageExec {
         ctx.session_config()
             .get_extension::<StageExec>()
             .ok_or(internal_datafusion_err!(
-                "ArrowFlightReadExec requires an ExecutionStage in the session config"
+                "missing ExecutionStage in session config"
             ))
     }
 
     pub fn child_stage(&self, i: usize) -> Result<&StageExec, DataFusionError> {
         self.child_stages_iter()
             .find(|s| s.num == i)
-            .ok_or(internal_datafusion_err!(
-                "ArrowFlightReadExec: no child stage with num {i}"
-            ))
+            .ok_or(internal_datafusion_err!("no child stage with num {i}"))
     }
 }
 
@@ -670,7 +668,7 @@ pub fn display_plan(
 /// An example of such a node would be:
 ///
 /// ```text
-///       ArrowFlightReadExec [label=<
+///       NetworkShuffleExec [label=<
 ///     <TABLE BORDER="0" CELLBORDER="0" CELLSPACING="0" CELLPADDING="0">
 ///         <TR>
 ///             <TD CELLBORDER="0">
@@ -686,7 +684,7 @@ pub fn display_plan(
 ///             <TD BORDER="0" CELLPADDING="0" CELLSPACING="0">
 ///                 <TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0">
 ///                     <TR>
-///                         <TD>ArrowFlightReadExec</TD>
+///                         <TD>NetworkShuffleExec</TD>
 ///                     </TR>
 ///                 </TABLE>
 ///             </TD>

tests/highly_distributed_query.rs

@@ -45,15 +45,15 @@ mod tests {
         assert_snapshot!(physical_distributed_str,
             @r"
         ┌───── Stage 4   Tasks: t0:[p0,p1,p2,p3,p4] 
-        │ ArrowFlightReadExec input_stage=3, input_partitions=5, input_tasks=1
+        │ NetworkShuffleExec input_stage=3, input_partitions=5, input_tasks=1
         └──────────────────────────────────────────────────
           ┌───── Stage 3   Tasks: t0:[p0,p1,p2,p3,p4] 
           │ RepartitionExec: partitioning=RoundRobinBatch(5), input_partitions=10
-          │   ArrowFlightReadExec input_stage=2, input_partitions=10, input_tasks=1
+          │   NetworkShuffleExec input_stage=2, input_partitions=10, input_tasks=1
           └──────────────────────────────────────────────────
             ┌───── Stage 2   Tasks: t0:[p0,p1,p2,p3,p4,p5,p6,p7,p8,p9] 
             │ RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
-            │   ArrowFlightReadExec input_stage=1, input_partitions=1, input_tasks=1
+            │   NetworkShuffleExec input_stage=1, input_partitions=1, input_tasks=1
             └──────────────────────────────────────────────────
               ┌───── Stage 1   Tasks: t0:[p0] 
               │ RepartitionExec: partitioning=RoundRobinBatch(1), input_partitions=1