[CARMEL-6283] DAGScheduler may stuck when there are too many result task and result is spilled (#1085)

Author: allenma

core/src/main/scala/org/apache/spark/scheduler/IterableJobWaiter.scala

@@ -49,6 +49,10 @@ private[spark] class IterableJobWaiter[U: ClassTag, R](
   private var spilledResultData: Option[Array[SpilledPartitionResult]] = None
   private val resultData: Array[U] = new Array[U](totalTasks)
 
+  // indicate whether the in memory result data has been cleaned after
+  // the result data is spilled to disk
+  @volatile private var inMemoryResultDataCleaned = false
+
   def jobFinished: Boolean = jobPromise.isCompleted
 
   def completionFuture: Future[Unit] = jobPromise.future
@@ -67,12 +71,16 @@ private[spark] class IterableJobWaiter[U: ClassTag, R](
       case spilledPartitionResult: Array[SpilledPartitionResult] =>
         spilledResultData = Some(spilledPartitionResult)
         // if spill is triggered for this task set, clean the cached in memory results
-        (0 until totalTasks).foreach { i =>
-          resultData(i) = null.asInstanceOf[U]
+        if (!inMemoryResultDataCleaned) {
+          cleanInMemResultData
+          inMemoryResultDataCleaned = true
         }
       case data: U =>
-        resultData(index) = data
-        resultStatCollector(index, result.asInstanceOf[U])
+        if (spilledResultData.isEmpty) {
+          // if the spilledResultData has been set, no need to save result to memory
+          resultData(index) = data
+          resultStatCollector(index, result.asInstanceOf[U])
+        }
 
       // case stats: MapOutputStatistics => // ignore
       case  _ =>
@@ -85,6 +93,12 @@ private[spark] class IterableJobWaiter[U: ClassTag, R](
     }
   }
 
+  def cleanInMemResultData: Unit = {
+    (0 until totalTasks).foreach { i =>
+      resultData(i) = null.asInstanceOf[U]
+    }
+  }
+
   override def jobFailed(exception: Exception): Unit = {
     if (!jobPromise.tryFailure(exception)) {
       logWarning("Ignore failure", exception)