sunchao commented on code in PR #55839:
URL: https://github.com/apache/spark/pull/55839#discussion_r3827573556
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -77,6 +78,17 @@ case class AdaptiveSparkPlanExec(
@transient private val lock = new Object()
+ // Access is serialized by the execution context's stage lifecycle lock.
+ @transient private val uncancelledObsoleteStageIds =
mutable.HashSet.empty[Int]
+
+ // Remember every local alias because each registers its own materialization
callback.
Review Comment:
Fixed in b783cdc3. Updated the comment to the suggested broader wording, so
it covers every local alias used by cancellation and failure handling,
including already-materialized reuse aliases.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -986,6 +1241,84 @@ case class AdaptiveExecutionContext(session:
SparkSession, qe: QueryExecution) {
val stageCache: TrieMap[SparkPlan, ExchangeQueryStageExec] =
new TrieMap[SparkPlan, ExchangeQueryStageExec]()
+ private val stageLifecycleLock = new Object
+
+ /**
+ * Serialize exchange-cache lookup, reuse, and cancellation reservations
across subqueries.
+ * Blocking cancellation and waits for a particular result happen outside
this query-wide lock.
+ */
+ private[adaptive] def withStageLifecycleLock[T](body: => T): T = {
+ stageLifecycleLock.synchronized(body)
+ }
+
+ /** Stage-scoped reservations keep a cancelling result cached without
blocking other stages. */
+ private val stageCancellationReservations =
+ new ConcurrentHashMap[AtomicReference[Option[Any]],
CompletableFuture[Unit]]()
+
+ /** Claim a result before waiting for its shuffle monitor or cancelling its
exchange. */
+ private[adaptive] def reserveStageCancellation(
+ resultOption: AtomicReference[Option[Any]]):
Option[CompletableFuture[Unit]] = {
+ val cancellation = new CompletableFuture[Unit]()
+ if (stageCancellationReservations.putIfAbsent(resultOption, cancellation)
== null) {
+ Some(cancellation)
+ } else {
+ None
+ }
+ }
+
+ /** Return the cancellation that a potential consumer of this particular
result must await. */
+ private[adaptive] def pendingStageCancellation(
+ resultOption: AtomicReference[Option[Any]]):
Option[CompletableFuture[Unit]] = {
+ Option(stageCancellationReservations.get(resultOption))
+ }
+
+ /** Publish the final cache state before allowing consumers of this result
to retry lookup. */
+ private[adaptive] def finishStageCancellation(
+ resultOption: AtomicReference[Option[Any]],
+ cancellation: CompletableFuture[Unit]): Unit = {
+ stageCancellationReservations.remove(resultOption, cancellation)
+ cancellation.complete(())
+ }
+
+ /**
+ * The adaptive plan that first cached each result. Identity, rather than
case-class equality,
+ * distinguishes independently planned subqueries with equivalent physical
input plans.
+ */
+ private val stageResultOwners =
+ new ConcurrentHashMap[AtomicReference[Option[Any]],
AdaptiveSparkPlanExec]()
+
+ /** Record the original owner atomically with exchange-cache insertion. */
+ private[adaptive] def registerStageOwner(
+ resultOption: AtomicReference[Option[Any]], owner:
AdaptiveSparkPlanExec): Unit = {
+ stageResultOwners.putIfAbsent(resultOption, owner)
+ }
+
+ /**
+ * Results reused by another adaptive plan in this execution context.
Cross-plan protection
+ * remains for the context's lifetime; aliases within one plan are not
considered shared.
+ */
+ private val sharedStageResults =
+ new ConcurrentHashMap[AtomicReference[Option[Any]], Boolean]()
+
+ /** Conservatively protect a result whose owner is unknown or explicitly
forced by a test. */
+ private[adaptive] def markSharedStageResult(resultOption:
AtomicReference[Option[Any]]): Unit = {
+ sharedStageResults.put(resultOption, true)
+ }
+
+ /** Protect a result only when another adaptive plan, rather than a local
alias, reuses it. */
+ private[adaptive] def markSharedStageResult(
+ resultOption: AtomicReference[Option[Any]], owner:
AdaptiveSparkPlanExec): Unit = {
+ val originalOwner = stageResultOwners.get(resultOption)
+ if (originalOwner == null || (originalOwner ne owner)) {
+ markSharedStageResult(resultOption)
+ }
+ }
+
+ /** Return whether an independently owned adaptive plan has reused this
result. */
Review Comment:
Fixed in b783cdc3. Reworded the one-argument marker, owner-aware overload,
and predicate documentation in terms of conservative protection, explicitly
covering unknown ownership, test-forced protection, and cross-plan reuse
without treating membership as proof of independent-plan reuse.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala:
##########
@@ -373,6 +376,1061 @@ class AdaptiveQueryExecSuite
}
}
+ test("non-empty global aggregate stage eliminates conditionless semi joins")
{
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2") {
+ val globalAggregate = spark.range(1).where("id < 0").repartition(2)
+ .agg(count("*").as("c"))
+ val df = testData.join(globalAggregate, Seq.empty[String], "left_semi")
+
+ val logicalJoin = df.queryExecution.optimizedPlan.collectFirst {
+ case join: Join => join
+ }.getOrElse(fail("expected a conditionless semi join before adaptive
execution"))
+ assert(logicalJoin.joinType == LeftSemi)
+ assert(logicalJoin.condition.isEmpty)
+
+ val initialPlan = df.queryExecution.executedPlan
+ .asInstanceOf[AdaptiveSparkPlanExec].initialPlan
+ assert(findTopLevelBaseJoin(initialPlan).size == 1)
+
+ val aggregate = collect(initialPlan) {
+ case aggregate: BaseAggregateExec if
aggregate.groupingExpressions.isEmpty => aggregate
+ }.headOption.getOrElse(fail("expected a global aggregate in the initial
adaptive plan"))
+ val emptyStage = TestExchangeQueryStageExec(
+ 0,
+ aggregate.child,
+ aggregate.child.canonicalized,
+ runtimeRowCount = Some(BigInt(0)))
+ emptyStage.resultOption.set(Some(()))
+ val aggregateStage = LogicalQueryStage(
+ logicalJoin.right, aggregate.withNewChildren(Seq(emptyStage)))
+ val rewrittenJoin = AQEPropagateEmptyRelation(logicalJoin.copy(right =
aggregateStage))
+ assert(rewrittenJoin.fastEquals(logicalJoin.left), rewrittenJoin)
+
+ checkAnswer(df, testData.collect().toSeq)
+
+ val finalPlan = stripAQEPlan(df.queryExecution.executedPlan)
+ assert(findTopLevelBaseJoin(finalPlan).isEmpty, finalPlan)
+ }
+ }
+
+ test("empty filtered global aggregate stage preserves its single output
row") {
Review Comment:
Fixed in b783cdc3. Renamed the test to "global aggregate over empty filtered
input preserves its single output row", distinguishing the empty filtered input
from the one-row aggregate result.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
