cloud-fan commented on code in PR #55839:
URL: https://github.com/apache/spark/pull/55839#discussion_r3826616233
##########
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:
**Nit:**
`stageIdsByResult` also records already-materialized reuse aliases, which do
not enter `newStages` and therefore do not register callbacks. Describe the
broader cancellation and failure-handling purpose instead.
```suggestion
// Remember every local alias so cancellation and failure handling cover
every stage ID.
```
##########
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:
**Nit:**
`isSharedStageResult` reports membership in `sharedStageResults`, not proof
of independent-plan reuse: the one-argument overload at
`AdaptiveSparkPlanExec.scala:1303` also inserts unknown-owner and test-forced
results. Please reword this comment and the one at line 1308 around
conservative protection; at line 1303, say the protection, rather than the
owner, can be test-forced.
##########
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:
**Nit:**
The filtered input is empty, while the aggregate stage produces one row.
Naming the input explicitly avoids describing the stage as both empty and
row-producing.
```suggestion
test("global aggregate over empty filtered input preserves its single
output row") {
```
--
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]
