From 0187cef0f284e6cb22cb3986c327c43304daf57d Mon Sep 17 00:00:00 2001
From: Matei Zaharia <matei@databricks.com>
Date: Mon, 24 Feb 2014 13:14:56 -0800
Subject: [PATCH] Fix removal from shuffleToMapStage to search for a key-value
pair with our stage instead of using our shuffleID.
---
.../main/scala/org/apache/spark/scheduler/DAGScheduler.scala | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
index 789d5e6699..dc5b25d845 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
@@ -381,8 +381,8 @@ class DAGScheduler(
running -= stage
}
stageToInfos -= stage
- for (shuffleDep <- stage.shuffleDep) {
- shuffleToMapStage.remove(shuffleDep.shuffleId)
+ for ((k, v) <- shuffleToMapStage.find(_._2 == stage)) {
+ shuffleToMapStage.remove(k)
}
if (pendingTasks.contains(stage) && !pendingTasks(stage).isEmpty) {
logDebug("Removing pending status for stage %d".format(stageId))
--
GitLab