From 0b055d6ee465248626438016fc56af64e003d9f9 Mon Sep 17 00:00:00 2001
From: Dylan Chen <zilin@singularity-data.com>
Date: Fri, 5 Jan 2024 17:55:46 +0800
Subject: [PATCH] fix

---
 src/frontend/src/scheduler/distributed/stage.rs | 7 ++++++-
 src/frontend/src/scheduler/local.rs             | 3 +++
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/src/frontend/src/scheduler/distributed/stage.rs b/src/frontend/src/scheduler/distributed/stage.rs
index 2b29cc5dc20a8..1335c9d1220d4 100644
--- a/src/frontend/src/scheduler/distributed/stage.rs
+++ b/src/frontend/src/scheduler/distributed/stage.rs
@@ -705,7 +705,9 @@ impl StageRunner {
                 .worker_node_manager
                 .manager
                 .get_workers_by_parallel_unit_ids(&parallel_unit_ids)?;
-
+            if candidates.is_empty() {
+                return Err(SchedulerError::EmptyWorkerNodes);
+            }
             return Ok(Some(
                 candidates
                     .get(self.stage.session_id.0 as usize % candidates.len())
@@ -736,6 +738,9 @@ impl StageRunner {
                 .worker_node_manager
                 .manager
                 .get_workers_by_parallel_unit_ids(&[pu])?;
+            if candidates.is_empty() {
+                return Err(SchedulerError::EmptyWorkerNodes);
+            }
             Ok(Some(candidates[0].clone()))
         } else {
             Ok(None)
diff --git a/src/frontend/src/scheduler/local.rs b/src/frontend/src/scheduler/local.rs
index e036967a54eef..e434c65bf97bd 100644
--- a/src/frontend/src/scheduler/local.rs
+++ b/src/frontend/src/scheduler/local.rs
@@ -580,6 +580,9 @@ impl LocalQueryExecution {
                     .worker_node_manager
                     .manager
                     .get_workers_by_parallel_unit_ids(&parallel_unit_ids)?;
+                if candidates.is_empty() {
+                    return Err(SchedulerError::EmptyWorkerNodes);
+                }
                 candidates
                     .get(stage.session_id.0 as usize % candidates.len())
                     .context("no available worker node for dml")?