From 4fa5b5f336ef0cca8597c830584ad267bdb53947 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=BC=A0=E6=9E=97=E4=BC=9F?= <lewiszlw520@gmail.com>
Date: Mon, 15 Jul 2024 22:30:01 +0800
Subject: [PATCH] Fix job hangs when partition count of plan is zero (#1024)

* Fix job hangs when partition count of plan is zero

* fix compilation issues

---------

Co-authored-by: Andy Grove <agrove@apache.org>
---
 ballista/core/src/serde/generated/ballista.rs |  1 -
 ballista/scheduler/src/state/mod.rs           | 24 +++++++++++++++++--
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/ballista/core/src/serde/generated/ballista.rs b/ballista/core/src/serde/generated/ballista.rs
index 50df950af..13a17f7ec 100644
--- a/ballista/core/src/serde/generated/ballista.rs
+++ b/ballista/core/src/serde/generated/ballista.rs
@@ -1,4 +1,3 @@
-// This file is @generated by prost-build.
 /// /////////////////////////////////////////////////////////////////////////////////////////////////
 /// Ballista Physical Plan
 /// /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/ballista/scheduler/src/state/mod.rs b/ballista/scheduler/src/state/mod.rs
index 90c9f6f31..b30c9d92e 100644
--- a/ballista/scheduler/src/state/mod.rs
+++ b/ballista/scheduler/src/state/mod.rs
@@ -15,10 +15,11 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
+use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
 use datafusion::datasource::listing::{ListingTable, ListingTableUrl};
 use datafusion::datasource::source_as_provider;
 use datafusion::error::DataFusionError;
+use datafusion::physical_plan::ExecutionPlanProperties;
 use std::any::type_name;
 use std::collections::HashMap;
 use std::sync::Arc;
@@ -39,6 +40,7 @@ use ballista_core::serde::protobuf::TaskStatus;
 use ballista_core::serde::BallistaCodec;
 use datafusion::logical_expr::LogicalPlan;
 use datafusion::physical_plan::display::DisplayableExecutionPlan;
+use datafusion::physical_plan::empty::EmptyExec;
 use datafusion::prelude::SessionContext;
 use datafusion_proto::logical_plan::AsLogicalPlan;
 use datafusion_proto::physical_plan::AsExecutionPlan;
@@ -408,8 +410,26 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerState<T,
             DisplayableExecutionPlan::new(plan.as_ref()).indent(false)
         );
 
+        let plan = plan.transform_down(&|node| {
+            if node.output_partitioning().partition_count() == 0 {
+                Ok(Transformed::yes(Arc::new(EmptyExec::new(node.schema()))))
+            } else {
+                Ok(Transformed::no(node))
+            }
+        })?;
+        debug!(
+            "Transformed physical plan: {}",
+            DisplayableExecutionPlan::new(plan.data.as_ref()).indent(false)
+        );
+
         self.task_manager
-            .submit_job(job_id, job_name, &session_ctx.session_id(), plan, queued_at)
+            .submit_job(
+                job_id,
+                job_name,
+                &session_ctx.session_id(),
+                plan.data,
+                queued_at,
+            )
             .await?;
 
         let elapsed = start.elapsed();