Skip to content

Commit 69b2ca8

Browse files
sunchaosunchao
committed
test: cover AQE DPP broadcast logical links
1 parent 0d82334 commit 69b2ca8

1 file changed

Lines changed: 150 additions & 3 deletions

File tree

‎spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala‎

Lines changed: 150 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,23 +19,29 @@
1919

2020
package org.apache.comet.rules
2121

22+
import java.util.concurrent.ConcurrentLinkedQueue
23+
24+
import scala.jdk.CollectionConverters._
2225
import scala.util.Random
2326

27+
import org.scalatest.PrivateMethodTester._
28+
2429
import org.apache.spark.sql._
2530
import org.apache.spark.sql.catalyst.FunctionIdentifier
2631
import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionInfo}
2732
import org.apache.spark.sql.catalyst.expressions.aggregate.{BloomFilterAggregate, Partial}
28-
import org.apache.spark.sql.catalyst.plans.logical.LocalRelation
33+
import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan}
34+
import org.apache.spark.sql.catalyst.rules.Rule
2935
import org.apache.spark.sql.comet._
3036
import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec
3137
import org.apache.spark.sql.execution._
32-
import org.apache.spark.sql.execution.adaptive.{LogicalQueryStage, QueryStageExec, ShuffleQueryStageExec}
38+
import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, BroadcastQueryStageExec, LogicalQueryStage, QueryStageExec, ShuffleQueryStageExec, SimpleCost, SimpleCostEvaluator}
3339
import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec}
3440
import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleExchangeExec}
3541
import org.apache.spark.sql.internal.SQLConf
3642
import org.apache.spark.sql.types.{DataTypes, StructField, StructType}
3743

38-
import org.apache.comet.{CometConf, CometExplainInfo, ExtendedExplainInfo}
44+
import org.apache.comet.{CometConf, CometExplainInfo, CometSparkSessionExtensions, ExtendedExplainInfo}
3945
import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus, isSpark42Plus}
4046
import org.apache.comet.serde.{Compatible, Unsupported}
4147
import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator}
@@ -47,6 +53,40 @@ import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator}
4753
*/
4854
class CometExecRuleSuite extends CometTestBase {
4955

56+
// The observers are active only during the DPP lifecycle regression below. AQE can prepare
57+
// subqueries on different threads, so publish the callbacks and pair each thread's invocations.
58+
@volatile private var beforeCometPreparation: SparkPlan => Unit = _ => ()
59+
@volatile private var afterCometPreparation: SparkPlan => Unit = _ => ()
60+
61+
override protected def createSparkSession: SparkSessionType = {
62+
SparkSession.clearActiveSession()
63+
SparkSession.clearDefaultSession()
64+
SparkSession
65+
.builder()
66+
.config(sparkContext.getConf)
67+
.withExtensions { extensions =>
68+
extensions.injectQueryStagePrepRule { _ =>
69+
new Rule[SparkPlan] {
70+
override def apply(plan: SparkPlan): SparkPlan = {
71+
beforeCometPreparation(plan)
72+
plan
73+
}
74+
}
75+
}
76+
new CometSparkSessionExtensions().apply(extensions)
77+
extensions.injectQueryStagePrepRule { _ =>
78+
new Rule[SparkPlan] {
79+
override def apply(plan: SparkPlan): SparkPlan = {
80+
afterCometPreparation(plan)
81+
plan
82+
}
83+
}
84+
}
85+
}
86+
.getOrCreate()
87+
.asInstanceOf[SparkSessionType]
88+
}
89+
5090
/** Helper method to apply CometExecRule and return the transformed plan */
5191
private def applyCometExecRule(plan: SparkPlan): SparkPlan = {
5292
CometExecRule(spark).apply(stripAQEPlan(plan))
@@ -183,6 +223,113 @@ class CometExecRuleSuite extends CometTestBase {
183223
}
184224
}
185225

226+
test("AQE DPP broadcast roots retain temporary logical links after an unchanged replan") {
227+
assume(isSpark35Plus, "Native AQE DPP requires Spark 3.5+")
228+
withSQLConf(
229+
SQLConf.USE_V1_SOURCE_LIST.key -> "parquet",
230+
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
231+
SQLConf.ADAPTIVE_FORCE_OPTIMIZE_SKEWED_JOIN.key -> "false",
232+
SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true",
233+
SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true",
234+
SQLConf.SHUFFLE_PARTITIONS.key -> "2",
235+
CometConf.COMET_SHUFFLE_MODE.key -> "native") {
236+
withTempDir { dir =>
237+
withTempView("dpp_link_fact", "dpp_link_dim") {
238+
withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
239+
spark
240+
.range(64)
241+
.selectExpr("CAST(id % 8 AS INT) AS k", "id AS v")
242+
.write
243+
.partitionBy("k")
244+
.parquet(s"$dir/fact")
245+
spark
246+
.range(32)
247+
.selectExpr(
248+
"CAST(id % 8 AS INT) AS k",
249+
"id AS v",
250+
"IF(id % 2 = 0, 'DE', 'US') AS country")
251+
.write
252+
.parquet(s"$dir/dim")
253+
}
254+
spark.read.parquet(s"$dir/fact").createOrReplaceTempView("dpp_link_fact")
255+
spark.read.parquet(s"$dir/dim").createOrReplaceTempView("dpp_link_dim")
256+
257+
assert(
258+
spark.sessionState.conf.getConf(SQLConf.ADAPTIVE_CUSTOM_COST_EVALUATOR_CLASS).isEmpty)
259+
type Replan = (SparkPlan, LogicalPlan)
260+
val pending = new ThreadLocal[List[Option[Replan]]] {
261+
override def initialValue(): List[Option[Replan]] = Nil
262+
}
263+
val observed = new ConcurrentLinkedQueue[(CometBroadcastExchangeExec, LogicalPlan)]()
264+
val tempTag = AdaptiveSparkPlanExec.TEMP_LOGICAL_PLAN_TAG
265+
val costEvaluator = SimpleCostEvaluator(forceOptimizeSkewedJoin = false)
266+
beforeCometPreparation = plan => {
267+
val replan = plan match {
268+
case broadcast: CometBroadcastExchangeExec =>
269+
broadcast.getTagValue(SparkPlan.LOGICAL_PLAN_TAG).collect {
270+
case stage: LogicalQueryStage =>
271+
assert(stage.physicalPlan eq broadcast)
272+
assert(broadcast.getTagValue(tempTag).exists(_ eq stage.logicalPlan))
273+
(broadcast.clone(), stage.logicalPlan)
274+
}
275+
case _ => None
276+
}
277+
pending.set(replan :: pending.get())
278+
}
279+
afterCometPreparation = plan => {
280+
val replan = pending.get().head
281+
val remaining = pending.get().tail
282+
if (remaining.isEmpty) pending.remove() else pending.set(remaining)
283+
replan.foreach { case (previous, logicalPlan) =>
284+
val broadcast = plan.asInstanceOf[CometBroadcastExchangeExec]
285+
assert(broadcast.logicalLink.exists(_ eq logicalPlan))
286+
assert(broadcast.getTagValue(tempTag).exists(_ eq logicalPlan))
287+
// Spark rejects an equal-cost candidate when its physical tree is unchanged.
288+
// Pin both inputs to that decision, including Comet's retained temporary link.
289+
assert(previous == broadcast)
290+
assert(costEvaluator.evaluateCost(previous) == SimpleCost(0))
291+
assert(costEvaluator.evaluateCost(broadcast) == SimpleCost(0))
292+
observed.add(
293+
(broadcast.clone().asInstanceOf[CometBroadcastExchangeExec], logicalPlan))
294+
}
295+
}
296+
try {
297+
val df = sql("""
298+
|SELECT /*+ BROADCAST(d) */ f.k, f.total, d.total
299+
|FROM (SELECT k, SUM(v) AS total FROM dpp_link_fact GROUP BY k) f
300+
|JOIN (SELECT k, SUM(v) AS total FROM dpp_link_dim
301+
| WHERE country = 'DE' GROUP BY k) d ON f.k = d.k
302+
|""".stripMargin)
303+
QueryTest.checkAnswer(
304+
df,
305+
(0 until 8 by 2).map(k => Row(k, 224L + 8L * k, 48L + 4L * k)),
306+
checkToRDD = false)
307+
val plan = df.queryExecution.executedPlan
308+
assert(collect(plan) { case b: CometBroadcastHashJoinExec => b }.nonEmpty)
309+
assert(collectWithSubqueries(plan) { case s: CometSubqueryBroadcastExec =>
310+
s
311+
}.nonEmpty)
312+
assert(!observed.isEmpty, "Expected a DPP broadcast root with a direct logical stage")
313+
observed.iterator().asScala.foreach { case (broadcast, logicalPlan) =>
314+
// Give the isolated snapshot conflicting links to pin Spark's TEMP-over-direct
315+
// precedence, which would otherwise be invisible after Comet repairs both.
316+
broadcast.setLogicalLink(LogicalQueryStage(logicalPlan, broadcast))
317+
val stage = BroadcastQueryStageExec(0, broadcast, broadcast.canonicalized)
318+
val setStageLink = PrivateMethod[Unit](Symbol("setLogicalLinkForNewQueryStage"))
319+
plan
320+
.asInstanceOf[AdaptiveSparkPlanExec]
321+
.invokePrivate(setStageLink(stage, broadcast))
322+
assert(stage.logicalLink.exists(_ eq logicalPlan))
323+
}
324+
} finally {
325+
beforeCometPreparation = _ => ()
326+
afterCometPreparation = _ => ()
327+
}
328+
}
329+
}
330+
}
331+
}
332+
186333
test("expression-level fallback reasons are rolled up onto the operator that falls back") {
187334
// Extended explain only walks plan nodes, so a reason recorded on a sub-expression is
188335
// invisible unless CometExecRule lifts it onto the enclosing operator. Disabling a single

0 commit comments

Comments
 (0)