feat: ShuffleReader shows upstream_stage

ballista/client/tests/context_checks.rs

@@ -1198,9 +1198,10 @@ mod supported {
             "|                   | ShuffleWriterExec: partitioning: None, metrics=[output_rows=..., input_rows=..., repart_time=..., write_time=...]                                                                                                                                                                                                                                                                                 |",
             "|                   |   ProjectionExec: expr=[count(Int64(1))@1 as count(*), id@0 as id], metrics=[output_rows=..., elapsed_compute=..., output_bytes=..., output_batches=..., expr_0_eval_time=..., expr_1_eval_time=...]                                                                                                                                                                                              |",
             "|                   |     AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[count(Int64(1))], metrics=[output_rows=..., elapsed_compute=..., output_bytes=..., output_batches=..., spill_count=..., spilled_bytes=..., spilled_rows=..., peak_mem_used=..., aggregate_arguments_time=..., aggregation_time=..., emitting_time=..., time_calculating_group_ids=...]                                          |",
-            "|                   |       ShuffleReaderExec: partitioning: Hash([id@0], 16), metrics=[output_rows=..., elapsed_compute=..., output_bytes=..., output_batches=...]                                                                                                                                                                                                                                                     |",
+            "|                   |       ShuffleReaderExec: upstream_stage: 1, partitioning: Hash([id@0], 16), metrics=[output_rows=..., elapsed_compute=..., output_bytes=..., output_batches=...]                                                                                                                                                                                                                                  |",
             "+-------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+",
         ];
+
         assert_batches_eq!(expected, &[sanitized]);
 
         Ok(())

ballista/core/src/execution_plans/shuffle_reader.rs

@@ -278,14 +278,14 @@ impl DisplayAs for ShuffleReaderExec {
                 if self.broadcast {
                     write!(
                         f,
-                        "ShuffleReaderExec: broadcast=true, upstream_partitions: {}",
-                        self.upstream_partition_count,
+                        "ShuffleReaderExec: upstream_stage: {}, broadcast: true, upstream_partition_count: {}",
+                        self.stage_id, self.upstream_partition_count,
                     )
                 } else {
                     write!(
                         f,
-                        "ShuffleReaderExec: partitioning: {}",
-                        self.properties.partitioning,
+                        "ShuffleReaderExec: upstream_stage: {}, partitioning: {}",
+                        self.stage_id, self.properties.partitioning,
                     )?;
                     if let Some(c) = &self.coalesce {
                         write!(
@@ -299,7 +299,8 @@ impl DisplayAs for ShuffleReaderExec {
                 }
             }
             DisplayFormatType::TreeRender => {
-                write!(f, "partitioning={}", self.properties.partitioning)
+                writeln!(f, "upstream_stage={}", self.stage_id)?;
+                writeln!(f, "partitioning={}", self.properties.partitioning)
             }
         }
     }

ballista/scheduler/src/state/aqe/test/alter_stages.rs

@@ -331,7 +331,7 @@ async fn should_support_cross_join() -> datafusion::error::Result<()> {
       ProjectionExec: expr=[big_col@1 as big_col, big_col@0 as big_col]
         CrossJoinExec
           CoalescePartitionsExec
-            ShuffleReaderExec: partitioning: UnknownPartitioning(2)
+            ShuffleReaderExec: upstream_stage: 0, partitioning: UnknownPartitioning(2)
           CooperativeExec
             StatisticsExec: col_count=1, row_count=Exact(262144)
     ");

ballista/scheduler/src/state/aqe/test/coalesce_rule.rs

@@ -427,7 +427,7 @@ async fn shuffle_reader_uses_coalesced_k_when_rule_fires() -> datafusion::error:
     ShuffleWriterExec: partitioning: None
       ProjectionExec: expr=[min(t.a)@1 as c0, max(t.b)@2 as c1, c@0 as c2]
         AggregateExec: mode=FinalPartitioned, gby=[c@0 as c], aggr=[min(t.a), max(t.b)]
-          ShuffleReaderExec: partitioning: Hash([c@0], 2), coalesce: 2 of 8
+          ShuffleReaderExec: upstream_stage: 0, partitioning: Hash([c@0], 2), coalesce: 2 of 8
     ");
 
     Ok(())

ballista/scheduler/src/state/aqe/test/plan_to_stages.rs

@@ -143,7 +143,7 @@ async fn should_split_plan_into_stages() -> datafusion::error::Result<()> {
     ShuffleWriterExec: partitioning: None
       ProjectionExec: expr=[min(t.a)@1 as c0, max(t.b)@2 as c1, c@0 as c2]
         AggregateExec: mode=FinalPartitioned, gby=[c@0 as c], aggr=[min(t.a), max(t.b)]
-          ShuffleReaderExec: partitioning: Hash([c@0], 2)
+          ShuffleReaderExec: upstream_stage: 0, partitioning: Hash([c@0], 2)
     ");
     planner.finalise_stage_internal(1, mock_partitions_with_statistics())?;