rebase

Author: summaryzb

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

@@ -51,11 +51,7 @@ private[spark] class JobWaiter[T](
    * all the tasks belonging to this job, it will fail this job with a SparkException.
    */
   def cancel(reason: Option[String] = None, quiet: Boolean = false): Unit = {
-    if (quiet) {
-      dagScheduler.cancelJob(jobId, reason, quiet)
-    } else {
-      dagScheduler.cancelJob(jobId, None, quiet)
-    }
+    dagScheduler.cancelJob(jobId, reason, quiet)
   }
 
   /**

sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala

@@ -384,6 +384,7 @@ case class AdaptiveSparkPlanExec(
                 }
               })
             }
+            stagesToCancel.clear()
           }
         }
         // Now that some stages have finished, we can try creating new stages.

sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala

@@ -78,9 +78,14 @@ trait ShuffleExchangeLike extends Exchange {
   private[sql] // Exposed for testing
   val futureAction = new AtomicReference[Option[FutureAction[MapOutputStatistics]]](None)
 
+  @volatile
   @transient
   private var isCancelled: Boolean = false
 
+  @volatile
+  @transient
+  private var quietly: Boolean = false
+
   @transient
   private lazy val triggerFuture: java.util.concurrent.Future[Any] = {
     SQLExecution.withThreadLocalCaptured(session, ShuffleExchangeExec.executionContext) {
@@ -90,7 +95,7 @@ trait ShuffleExchangeLike extends Exchange {
         executeQuery(null)
         // Submit shuffle job if not cancelled.
         this.synchronized {
-          if (isCancelled) {
+          if (isCancelled && !quietly) {
             promise.tryFailure(new SparkException("Shuffle cancelled."))
           } else {
             val shuffleJob = RDDOperationScope.withScope(sparkContext, nodeName, false, true) {
@@ -125,9 +130,15 @@ trait ShuffleExchangeLike extends Exchange {
    * Cancels the shuffle job with an optional reason.
    */
   final def cancelShuffleJob(reason: Option[String], quiet: Boolean): Unit = this.synchronized {
-    if (!isCancelled) {
-      isCancelled = true
-      futureAction.get().foreach(_.cancel(reason, quiet))
+    this.synchronized {
+      if (!isCancelled) {
+        isCancelled = true
+        if (quiet) {
+          quietly = quiet
+          promise.tryFailure(new SparkAQEStageCancelException)
+        }
+        futureAction.get().foreach(_.cancel(reason, quiet))
+      }
     }
   }
 

sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala

@@ -255,44 +255,6 @@ class AdaptiveQueryExecSuite
     }
   }
 
-  test("SPARK-52024: Support AQE Cancel Empty") {
-    // make the non-empty side compute more time
-    // so that empty side complete first to trigger the rule
-    spark.udf.register("fake_udf", (input: Int) => {
-      Thread.sleep(20)
-      input
-    })
-    var withCancelEmpty = -1L
-    var withoutCancelEmpty = -1L
-    Seq(true, false).foreach { enable =>
-      withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
-        SQLConf.ADAPTIVE_EMPTY_TRIGGER_CANCEL_ENABLED.key -> String.valueOf(enable)) {
-        val query =
-          """
-            |SELECT t.key1
-            |FROM emptyTestData join (SELECT testData.key as key1
-            |FROM testData join testData2 ON fake_udf(testData.key)=fake_udf(testData2.a) ) t
-            |on t.key1 = emptyTestData.key
-            |union
-            |SELECT testData.key
-            |FROM testData join testData2 ON testData.key=testData2.a
-            |""".stripMargin
-        runAdaptiveAndVerifyResult(query)
-
-        val start = System.currentTimeMillis()
-        sql(query).collect()
-        val end = System.currentTimeMillis()
-        if (enable) {
-          withCancelEmpty = end - start
-        } else {
-          withoutCancelEmpty = end - start
-        }
-      }
-    }
-    assert(withCancelEmpty * 2 < withoutCancelEmpty, "withCancelEmpty should be less")
-  }
-
-
   test("Reuse the parallelism of coalesced shuffle in local shuffle read") {
     withSQLConf(
       SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",