Fix index out of bounds panic (#819)

yahoNanJing · kyotoYaho · web-flow · commit 553b9a7d74b6 · 2023-06-28T09:26:57.000+08:00
* Remove output_partitioning from TaskDescription

* Remove output_partitioning from execution stages

* Remove input_partition_count from UnresolvedShuffleExec

* Add input stage id for ShuffleReaderExec

* Correct the behavior of output_partitioning() for ShuffleWriterExec

---------

Co-authored-by: yangzhong &lt;yangzhong@ebay.com&gt;
diff --git a/ballista/core/proto/ballista.proto b/ballista/core/proto/ballista.proto
@@ -49,13 +49,14 @@ message ShuffleWriterExecNode {
 message UnresolvedShuffleExecNode {
   uint32 stage_id = 1;
   datafusion.Schema schema = 2;
-  uint32 input_partition_count = 3;
   uint32 output_partition_count = 4;
 }
 
 message ShuffleReaderExecNode {
   repeated ShuffleReaderPartition partition = 1;
   datafusion.Schema schema = 2;
+  // The stage to read from
+  uint32 stage_id = 3;
 }
 
 message ShuffleReaderPartition {
@@ -98,7 +99,6 @@ message ExecutionGraphStage {
 
 message UnResolvedStage {
   uint32 stage_id = 1;
-  datafusion.PhysicalHashRepartition output_partitioning = 2;
   repeated uint32 output_links = 3;
   repeated  GraphStageInput inputs = 4;
   bytes plan = 5;
@@ -109,7 +109,6 @@ message UnResolvedStage {
 message ResolvedStage {
   uint32 stage_id = 1;
   uint32 partitions = 2;
-  datafusion.PhysicalHashRepartition output_partitioning = 3;
   repeated uint32 output_links = 4;
   repeated  GraphStageInput inputs = 5;
   bytes plan = 6;
@@ -120,7 +119,6 @@ message ResolvedStage {
 message SuccessfulStage {
   uint32 stage_id = 1;
   uint32 partitions = 2;
-  datafusion.PhysicalHashRepartition output_partitioning = 3;
   repeated uint32 output_links = 4;
   repeated  GraphStageInput inputs = 5;
   bytes plan = 6;
@@ -132,7 +130,6 @@ message SuccessfulStage {
 message FailedStage {
   uint32 stage_id = 1;
   uint32 partitions = 2;
-  datafusion.PhysicalHashRepartition output_partitioning = 3;
   repeated uint32 output_links = 4;
   bytes plan = 5;
   repeated TaskInfo task_infos = 6;
diff --git a/ballista/core/src/execution_plans/shuffle_reader.rs b/ballista/core/src/execution_plans/shuffle_reader.rs
@@ -58,22 +58,26 @@ use tokio_stream::wrappers::ReceiverStream;
 /// being executed by an executor
 #[derive(Debug, Clone)]
 pub struct ShuffleReaderExec {
+    /// The query stage id to read from
+    pub stage_id: usize,
+    pub(crate) schema: SchemaRef,
     /// Each partition of a shuffle can read data from multiple locations
     pub partition: Vec<Vec<PartitionLocation>>,
-    pub(crate) schema: SchemaRef,
     /// Execution metrics
     metrics: ExecutionPlanMetricsSet,
 }
 
 impl ShuffleReaderExec {
     /// Create a new ShuffleReaderExec
     pub fn try_new(
+        stage_id: usize,
         partition: Vec<Vec<PartitionLocation>>,
         schema: SchemaRef,
     ) -> Result<Self> {
         Ok(Self {
-            partition,
+            stage_id,
             schema,
+            partition,
             metrics: ExecutionPlanMetricsSet::new(),
         })
     }
@@ -513,13 +517,14 @@ mod tests {
         ]);
 
         let job_id = "test_job_1";
+        let input_stage_id = 2;
         let mut partitions: Vec<PartitionLocation> = vec![];
         for partition_id in 0..4 {
             partitions.push(PartitionLocation {
                 map_partition_id: 0,
                 partition_id: PartitionId {
                     job_id: job_id.to_string(),
-                    stage_id: 2,
+                    stage_id: input_stage_id,
                     partition_id,
                 },
                 executor_meta: ExecutorMetadata {
@@ -534,8 +539,11 @@ mod tests {
             })
         }
 
-        let shuffle_reader_exec =
-            ShuffleReaderExec::try_new(vec![partitions], Arc::new(schema))?;
+        let shuffle_reader_exec = ShuffleReaderExec::try_new(
+            input_stage_id,
+            vec![partitions],
+            Arc::new(schema),
+        )?;
         let mut stream = shuffle_reader_exec.execute(0, task_ctx)?;
         let batches = utils::collect_stream(&mut stream).await;
 
diff --git a/ballista/core/src/execution_plans/shuffle_writer.rs b/ballista/core/src/execution_plans/shuffle_writer.rs
@@ -71,7 +71,8 @@ pub struct ShuffleWriterExec {
     plan: Arc<dyn ExecutionPlan>,
     /// Path to write output streams to
     work_dir: String,
-    /// Optional shuffle output partitioning
+    /// Optional shuffle output partitioning.
+    /// If it's none, it means there's no need to do repartitioning.
     shuffle_output_partitioning: Option<Partitioning>,
     /// Execution metrics
     metrics: ExecutionPlanMetricsSet,
@@ -134,6 +135,11 @@ impl ShuffleWriterExec {
         self.stage_id
     }
 
+    /// Get the input partition count
+    pub fn input_partition_count(&self) -> usize {
+        self.plan.output_partitioning().partition_count()
+    }
+
     /// Get the true output partitioning
     pub fn shuffle_output_partitioning(&self) -> Option<&Partitioning> {
         self.shuffle_output_partitioning.as_ref()
@@ -297,12 +303,12 @@ impl ExecutionPlan for ShuffleWriterExec {
         self.plan.schema()
     }
 
+    /// If [`shuffle_output_partitioning`] is none, then there's no need to do repartitioning.
+    /// Therefore, the partition is the same as its input plan's.
     fn output_partitioning(&self) -> Partitioning {
-        // This operator needs to be executed once for each *input* partition and there
-        // isn't really a mechanism yet in DataFusion to support this use case so we report
-        // the input partitioning as the output partitioning here. The executor reports
-        // output partition meta data back to the scheduler.
-        self.plan.output_partitioning()
+        self.shuffle_output_partitioning
+            .clone()
+            .unwrap_or_else(|| self.plan.output_partitioning())
     }
 
     fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
diff --git a/ballista/core/src/execution_plans/unresolved_shuffle.rs b/ballista/core/src/execution_plans/unresolved_shuffle.rs
@@ -38,9 +38,6 @@ pub struct UnresolvedShuffleExec {
     // The schema this node will have once it is replaced with a ShuffleReaderExec
     pub schema: SchemaRef,
 
-    // The number of shuffle writer partition tasks that will produce the partitions
-    pub input_partition_count: usize,
-
     // The partition count this node will have once it is replaced with a ShuffleReaderExec
     pub output_partition_count: usize,
 }
@@ -50,13 +47,11 @@ impl UnresolvedShuffleExec {
     pub fn new(
         stage_id: usize,
         schema: SchemaRef,
-        input_partition_count: usize,
         output_partition_count: usize,
     ) -> Self {
         Self {
             stage_id,
             schema,
-            input_partition_count,
             output_partition_count,
         }
     }
diff --git a/ballista/core/src/serde/generated/ballista.rs b/ballista/core/src/serde/generated/ballista.rs
@@ -45,8 +45,6 @@ pub struct UnresolvedShuffleExecNode {
     pub stage_id: u32,
     #[prost(message, optional, tag = "2")]
     pub schema: ::core::option::Option<::datafusion_proto::protobuf::Schema>,
-    #[prost(uint32, tag = "3")]
-    pub input_partition_count: u32,
     #[prost(uint32, tag = "4")]
     pub output_partition_count: u32,
 }
@@ -57,6 +55,9 @@ pub struct ShuffleReaderExecNode {
     pub partition: ::prost::alloc::vec::Vec<ShuffleReaderPartition>,
     #[prost(message, optional, tag = "2")]
     pub schema: ::core::option::Option<::datafusion_proto::protobuf::Schema>,
+    /// The stage to read from
+    #[prost(uint32, tag = "3")]
+    pub stage_id: u32,
 }
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]
@@ -132,10 +133,6 @@ pub mod execution_graph_stage {
 pub struct UnResolvedStage {
     #[prost(uint32, tag = "1")]
     pub stage_id: u32,
-    #[prost(message, optional, tag = "2")]
-    pub output_partitioning: ::core::option::Option<
-        ::datafusion_proto::protobuf::PhysicalHashRepartition,
-    >,
     #[prost(uint32, repeated, tag = "3")]
     pub output_links: ::prost::alloc::vec::Vec<u32>,
     #[prost(message, repeated, tag = "4")]
@@ -156,10 +153,6 @@ pub struct ResolvedStage {
     pub stage_id: u32,
     #[prost(uint32, tag = "2")]
     pub partitions: u32,
-    #[prost(message, optional, tag = "3")]
-    pub output_partitioning: ::core::option::Option<
-        ::datafusion_proto::protobuf::PhysicalHashRepartition,
-    >,
     #[prost(uint32, repeated, tag = "4")]
     pub output_links: ::prost::alloc::vec::Vec<u32>,
     #[prost(message, repeated, tag = "5")]
@@ -180,10 +173,6 @@ pub struct SuccessfulStage {
     pub stage_id: u32,
     #[prost(uint32, tag = "2")]
     pub partitions: u32,
-    #[prost(message, optional, tag = "3")]
-    pub output_partitioning: ::core::option::Option<
-        ::datafusion_proto::protobuf::PhysicalHashRepartition,
-    >,
     #[prost(uint32, repeated, tag = "4")]
     pub output_links: ::prost::alloc::vec::Vec<u32>,
     #[prost(message, repeated, tag = "5")]
@@ -204,10 +193,6 @@ pub struct FailedStage {
     pub stage_id: u32,
     #[prost(uint32, tag = "2")]
     pub partitions: u32,
-    #[prost(message, optional, tag = "3")]
-    pub output_partitioning: ::core::option::Option<
-        ::datafusion_proto::protobuf::PhysicalHashRepartition,
-    >,
     #[prost(uint32, repeated, tag = "4")]
     pub output_links: ::prost::alloc::vec::Vec<u32>,
     #[prost(bytes = "vec", tag = "5")]
diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs
@@ -157,6 +157,7 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec {
                 )?))
             }
             PhysicalPlanType::ShuffleReader(shuffle_reader) => {
+                let stage_id = shuffle_reader.stage_id as usize;
                 let schema = Arc::new(convert_required!(shuffle_reader.schema)?);
                 let partition_location: Vec<Vec<PartitionLocation>> = shuffle_reader
                     .partition
@@ -175,16 +176,14 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec {
                     })
                     .collect::<Result<Vec<_>, DataFusionError>>()?;
                 let shuffle_reader =
-                    ShuffleReaderExec::try_new(partition_location, schema)?;
+                    ShuffleReaderExec::try_new(stage_id, partition_location, schema)?;
                 Ok(Arc::new(shuffle_reader))
             }
             PhysicalPlanType::UnresolvedShuffle(unresolved_shuffle) => {
                 let schema = Arc::new(convert_required!(unresolved_shuffle.schema)?);
                 Ok(Arc::new(UnresolvedShuffleExec {
                     stage_id: unresolved_shuffle.stage_id as usize,
                     schema,
-                    input_partition_count: unresolved_shuffle.input_partition_count
-                        as usize,
                     output_partition_count: unresolved_shuffle.output_partition_count
                         as usize,
                 }))
@@ -237,6 +236,7 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec {
 
             Ok(())
         } else if let Some(exec) = node.as_any().downcast_ref::<ShuffleReaderExec>() {
+            let stage_id = exec.stage_id as u32;
             let mut partition = vec![];
             for location in &exec.partition {
                 partition.push(protobuf::ShuffleReaderPartition {
@@ -255,6 +255,7 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec {
             let proto = protobuf::BallistaPhysicalPlanNode {
                 physical_plan_type: Some(PhysicalPlanType::ShuffleReader(
                     protobuf::ShuffleReaderExecNode {
+                        stage_id,
                         partition,
                         schema: Some(exec.schema().as_ref().try_into()?),
                     },
@@ -273,7 +274,6 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec {
                     protobuf::UnresolvedShuffleExecNode {
                         stage_id: exec.stage_id as u32,
                         schema: Some(exec.schema().as_ref().try_into()?),
-                        input_partition_count: exec.input_partition_count as u32,
                         output_partition_count: exec.output_partition_count as u32,
                     },
                 )),
diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs
@@ -178,10 +178,6 @@ fn create_unresolved_shuffle(
         shuffle_writer.stage_id(),
         shuffle_writer.schema(),
         shuffle_writer.output_partitioning().partition_count(),
-        shuffle_writer
-            .shuffle_output_partitioning()
-            .map(|p| p.partition_count())
-            .unwrap_or_else(|| shuffle_writer.output_partitioning().partition_count()),
     ))
 }
 
@@ -246,6 +242,7 @@ pub fn remove_unresolved_shuffles(
                     .join("\n")
             );
             new_children.push(Arc::new(ShuffleReaderExec::try_new(
+                unresolved_shuffle.stage_id,
                 relevant_locations,
                 unresolved_shuffle.schema().clone(),
             )?))
@@ -265,15 +262,13 @@ pub fn rollback_resolved_shuffles(
     let mut new_children: Vec<Arc<dyn ExecutionPlan>> = vec![];
     for child in stage.children() {
         if let Some(shuffle_reader) = child.as_any().downcast_ref::<ShuffleReaderExec>() {
-            let partition_locations = &shuffle_reader.partition;
-            let output_partition_count = partition_locations.len();
-            let input_partition_count = partition_locations[0].len();
-            let stage_id = partition_locations[0][0].partition_id.stage_id;
+            let output_partition_count =
+                shuffle_reader.output_partitioning().partition_count();
+            let stage_id = shuffle_reader.stage_id;
 
             let unresolved_shuffle = Arc::new(UnresolvedShuffleExec::new(
                 stage_id,
                 shuffle_reader.schema(),
-                input_partition_count,
                 output_partition_count,
             ));
             new_children.push(unresolved_shuffle);
@@ -392,7 +387,6 @@ mod test {
         let unresolved_shuffle =
             downcast_exec!(unresolved_shuffle, UnresolvedShuffleExec);
         assert_eq!(unresolved_shuffle.stage_id, 1);
-        assert_eq!(unresolved_shuffle.input_partition_count, 2);
         assert_eq!(unresolved_shuffle.output_partition_count, 2);
 
         // verify stage 2
@@ -402,7 +396,6 @@ mod test {
         let unresolved_shuffle =
             downcast_exec!(unresolved_shuffle, UnresolvedShuffleExec);
         assert_eq!(unresolved_shuffle.stage_id, 2);
-        assert_eq!(unresolved_shuffle.input_partition_count, 2);
         assert_eq!(unresolved_shuffle.output_partition_count, 2);
 
         Ok(())
@@ -555,15 +548,13 @@ order by
         let join_input_1 = join_input_1.children()[0].clone();
         let unresolved_shuffle_reader_1 =
             downcast_exec!(join_input_1, UnresolvedShuffleExec);
-        assert_eq!(unresolved_shuffle_reader_1.input_partition_count, 2); // lineitem
         assert_eq!(unresolved_shuffle_reader_1.output_partition_count, 2);
 
         let join_input_2 = join.children()[1].clone();
         // skip CoalesceBatches
         let join_input_2 = join_input_2.children()[0].clone();
         let unresolved_shuffle_reader_2 =
             downcast_exec!(join_input_2, UnresolvedShuffleExec);
-        assert_eq!(unresolved_shuffle_reader_2.input_partition_count, 1); // orders
         assert_eq!(unresolved_shuffle_reader_2.output_partition_count, 2);
 
         // final partitioned hash aggregate
diff --git a/ballista/scheduler/src/scheduler_server/mod.rs b/ballista/scheduler/src/scheduler_server/mod.rs
@@ -460,10 +460,7 @@ mod test {
             if let Some(task) = task {
                 let mut partitions: Vec<ShuffleWritePartition> = vec![];
 
-                let num_partitions = task
-                    .output_partitioning
-                    .map(|p| p.partition_count())
-                    .unwrap_or(1);
+                let num_partitions = task.get_output_partition_number();
 
                 for partition_id in 0..num_partitions {
                     partitions.push(ShuffleWritePartition {
diff --git a/ballista/scheduler/src/state/execution_graph.rs b/ballista/scheduler/src/state/execution_graph.rs
diff --git a/ballista/scheduler/src/state/execution_graph/execution_stage.rs b/ballista/scheduler/src/state/execution_graph/execution_stage.rs
diff --git a/ballista/scheduler/src/state/execution_graph_dot.rs b/ballista/scheduler/src/state/execution_graph_dot.rs
diff --git a/ballista/scheduler/src/test_utils.rs b/ballista/scheduler/src/test_utils.rs
