Skip to content

fix: restore localReads assertion in AdaptiveQueryExecSuite diffs #6382

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mizulun wants to merge 3 commits into
base: main
Choose a base branch
from
Open

fix: restore localReads assertion in AdaptiveQueryExecSuite diffs #6382

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7ee3f02
fix: restore localReads assertion in AdaptiveQueryExecSuite diffs
mizulun Sep 29, 2026
f39539f
Merge branch 'main' into 6122-v1
mizulun Sep 29, 2026
b8c703d
fix: restore localReads assertion in Spark 3.4.3 diff
mizulun Oct 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 30 additions & 40 deletions dev/diffs/3.4.3.diff
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1563,7 +1563,7 @@ index ac710c32296..2854b433dd3 100644

import testImplicits._
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
index 593bd7bb4ba..b327d84d5cc 100644
index 593bd7bb4ba..9467f497cb9 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
@@ -26,9 +26,11 @@ import org.scalatest.time.SpanSugar._
Expand Down Expand Up @@ -1674,17 +1674,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80",
@@ -273,7 +289,8 @@ class AdaptiveQueryExecSuite
val localReads = collect(adaptivePlan) {
case read: AQEShuffleReadExec if read.isLocalRead => read
}
- assert(localReads.length == 2)
+ // Comet shuffle changes shuffle metrics
+ assert(localReads.length == 1)
val localShuffleRDD0 = localReads(0).execute().asInstanceOf[ShuffledRowRDD]
val localShuffleRDD1 = localReads(1).execute().asInstanceOf[ShuffledRowRDD]
// the final parallelism is math.max(1, numReduces / numMappers): math.max(1, 5/2) = 2
@@ -298,7 +315,9 @@ class AdaptiveQueryExecSuite
@@ -298,7 +314,9 @@ class AdaptiveQueryExecSuite
.groupBy($"a").count()
checkAnswer(testDf, Seq())
val plan = testDf.queryExecution.executedPlan
Expand All @@ -1695,7 +1685,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
val coalescedReads = collect(plan) {
case r: AQEShuffleReadExec => r
}
@@ -312,7 +331,9 @@ class AdaptiveQueryExecSuite
@@ -312,7 +330,9 @@ class AdaptiveQueryExecSuite
.groupBy($"a").count()
checkAnswer(testDf, Seq())
val plan = testDf.queryExecution.executedPlan
Expand All @@ -1706,7 +1696,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
val coalescedReads = collect(plan) {
case r: AQEShuffleReadExec => r
}
@@ -322,7 +343,7 @@ class AdaptiveQueryExecSuite
@@ -322,7 +342,7 @@ class AdaptiveQueryExecSuite
}
}

Expand All @@ -1715,7 +1705,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
@@ -337,7 +358,7 @@ class AdaptiveQueryExecSuite
@@ -337,7 +357,7 @@ class AdaptiveQueryExecSuite
}
}

Expand All @@ -1724,7 +1714,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
@@ -353,7 +374,7 @@ class AdaptiveQueryExecSuite
@@ -353,7 +373,7 @@ class AdaptiveQueryExecSuite
}
}

Expand All @@ -1733,7 +1723,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
@@ -398,7 +419,7 @@ class AdaptiveQueryExecSuite
@@ -398,7 +418,7 @@ class AdaptiveQueryExecSuite
}
}

Expand All @@ -1742,7 +1732,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
@@ -443,7 +464,7 @@ class AdaptiveQueryExecSuite
@@ -443,7 +463,7 @@ class AdaptiveQueryExecSuite
}
}

Expand All @@ -1751,7 +1741,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "500") {
@@ -489,7 +510,7 @@ class AdaptiveQueryExecSuite
@@ -489,7 +509,7 @@ class AdaptiveQueryExecSuite
}
}

Expand All @@ -1760,7 +1750,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
@@ -508,7 +529,7 @@ class AdaptiveQueryExecSuite
@@ -508,7 +528,7 @@ class AdaptiveQueryExecSuite
}
}

Expand All @@ -1769,7 +1759,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
@@ -539,7 +560,9 @@ class AdaptiveQueryExecSuite
@@ -539,7 +559,9 @@ class AdaptiveQueryExecSuite
assert(smj.size == 1)
val bhj = findTopLevelBroadcastHashJoin(adaptivePlan)
assert(bhj.size == 1)
Expand All @@ -1780,7 +1770,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
// Even with local shuffle read, the query stage reuse can also work.
val ex = findReusedExchange(adaptivePlan)
assert(ex.nonEmpty)
@@ -560,7 +583,9 @@ class AdaptiveQueryExecSuite
@@ -560,7 +582,9 @@ class AdaptiveQueryExecSuite
assert(smj.size == 1)
val bhj = findTopLevelBroadcastHashJoin(adaptivePlan)
assert(bhj.size == 1)
Expand All @@ -1791,7 +1781,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
// Even with local shuffle read, the query stage reuse can also work.
val ex = findReusedExchange(adaptivePlan)
assert(ex.isEmpty)
@@ -569,7 +594,8 @@ class AdaptiveQueryExecSuite
@@ -569,7 +593,8 @@ class AdaptiveQueryExecSuite
}
}

Expand All @@ -1801,7 +1791,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "20000000",
@@ -664,7 +690,8 @@ class AdaptiveQueryExecSuite
@@ -664,7 +689,8 @@ class AdaptiveQueryExecSuite
val bhj = findTopLevelBroadcastHashJoin(adaptivePlan)
assert(bhj.size == 1)
// There is still a SMJ, and its two shuffles can't apply local read.
Expand All @@ -1811,7 +1801,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
}
}

@@ -786,7 +813,8 @@ class AdaptiveQueryExecSuite
@@ -786,7 +812,8 @@ class AdaptiveQueryExecSuite
}
}

Expand All @@ -1821,7 +1811,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
Seq("SHUFFLE_MERGE", "SHUFFLE_HASH").foreach { joinHint =>
def getJoinNode(plan: SparkPlan): Seq[ShuffledJoin] = if (joinHint == "SHUFFLE_MERGE") {
findTopLevelSortMergeJoin(plan)
@@ -1004,7 +1032,8 @@ class AdaptiveQueryExecSuite
@@ -1004,7 +1031,8 @@ class AdaptiveQueryExecSuite
}
}

Expand All @@ -1831,7 +1821,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") {
val (_, adaptivePlan) = runAdaptiveAndVerifyResult(
"SELECT key FROM testData GROUP BY key")
@@ -1599,7 +1628,7 @@ class AdaptiveQueryExecSuite
@@ -1599,7 +1627,7 @@ class AdaptiveQueryExecSuite
val (_, adaptivePlan) = runAdaptiveAndVerifyResult(
"SELECT id FROM v1 GROUP BY id DISTRIBUTE BY id")
assert(collect(adaptivePlan) {
Expand All @@ -1840,7 +1830,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
}.length == 1)
}
}
@@ -1679,7 +1708,8 @@ class AdaptiveQueryExecSuite
@@ -1679,7 +1707,8 @@ class AdaptiveQueryExecSuite
}
}

Expand All @@ -1850,7 +1840,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
def hasRepartitionShuffle(plan: SparkPlan): Boolean = {
find(plan) {
case s: ShuffleExchangeLike =>
@@ -1864,6 +1894,9 @@ class AdaptiveQueryExecSuite
@@ -1864,6 +1893,9 @@ class AdaptiveQueryExecSuite
def checkNoCoalescePartitions(ds: Dataset[Row], origin: ShuffleOrigin): Unit = {
assert(collect(ds.queryExecution.executedPlan) {
case s: ShuffleExchangeExec if s.shuffleOrigin == origin && s.numPartitions == 2 => s
Expand All @@ -1860,7 +1850,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
}.size == 1)
ds.collect()
val plan = ds.queryExecution.executedPlan
@@ -1872,6 +1905,9 @@ class AdaptiveQueryExecSuite
@@ -1872,6 +1904,9 @@ class AdaptiveQueryExecSuite
}.isEmpty)
assert(collect(plan) {
case s: ShuffleExchangeExec if s.shuffleOrigin == origin && s.numPartitions == 2 => s
Expand All @@ -1870,7 +1860,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
}.size == 1)
checkAnswer(ds, testData)
}
@@ -2028,7 +2064,8 @@ class AdaptiveQueryExecSuite
@@ -2028,7 +2063,8 @@ class AdaptiveQueryExecSuite
}
}

Expand All @@ -1880,7 +1870,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
withTempView("t1", "t2") {
def checkJoinStrategy(shouldShuffleHashJoin: Boolean): Unit = {
Seq("100", "100000").foreach { size =>
@@ -2114,7 +2151,8 @@ class AdaptiveQueryExecSuite
@@ -2114,7 +2150,8 @@ class AdaptiveQueryExecSuite
}
}

Expand All @@ -1890,7 +1880,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
withTempView("v") {
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
@@ -2213,7 +2251,7 @@ class AdaptiveQueryExecSuite
@@ -2213,7 +2250,7 @@ class AdaptiveQueryExecSuite
runAdaptiveAndVerifyResult(s"SELECT $repartition key1 FROM skewData1 " +
s"JOIN skewData2 ON key1 = key2 GROUP BY key1")
val shuffles1 = collect(adaptive1) {
Expand All @@ -1899,7 +1889,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
}
assert(shuffles1.size == 3)
// shuffles1.head is the top-level shuffle under the Aggregate operator
@@ -2226,7 +2264,7 @@ class AdaptiveQueryExecSuite
@@ -2226,7 +2263,7 @@ class AdaptiveQueryExecSuite
runAdaptiveAndVerifyResult(s"SELECT $repartition key1 FROM skewData1 " +
s"JOIN skewData2 ON key1 = key2")
val shuffles2 = collect(adaptive2) {
Expand All @@ -1908,7 +1898,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
}
if (hasRequiredDistribution) {
assert(shuffles2.size == 3)
@@ -2260,7 +2298,8 @@ class AdaptiveQueryExecSuite
@@ -2260,7 +2297,8 @@ class AdaptiveQueryExecSuite
}
}

Expand All @@ -1918,15 +1908,15 @@ index 593bd7bb4ba..b327d84d5cc 100644
CostEvaluator.instantiate(
classOf[SimpleShuffleSortCostEvaluator].getCanonicalName, spark.sparkContext.getConf)
intercept[IllegalArgumentException] {
@@ -2404,6 +2443,7 @@ class AdaptiveQueryExecSuite
@@ -2404,6 +2442,7 @@ class AdaptiveQueryExecSuite
val (_, adaptive) = runAdaptiveAndVerifyResult(query)
assert(adaptive.collect {
case sort: SortExec => sort
+ case sort: CometSortExec => sort
}.size == 1)
val read = collect(adaptive) {
case read: AQEShuffleReadExec => read
@@ -2421,7 +2461,8 @@ class AdaptiveQueryExecSuite
@@ -2421,7 +2460,8 @@ class AdaptiveQueryExecSuite
}
}

Expand All @@ -1936,7 +1926,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
withTempView("v") {
withSQLConf(
SQLConf.ADAPTIVE_OPTIMIZE_SKEWS_IN_REBALANCE_PARTITIONS_ENABLED.key -> "true",
@@ -2533,7 +2574,7 @@ class AdaptiveQueryExecSuite
@@ -2533,7 +2573,7 @@ class AdaptiveQueryExecSuite
runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " +
"JOIN skewData3 ON value2 = value3")
val shuffles1 = collect(adaptive1) {
Expand All @@ -1945,7 +1935,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
}
assert(shuffles1.size == 4)
val smj1 = findTopLevelSortMergeJoin(adaptive1)
@@ -2544,7 +2585,7 @@ class AdaptiveQueryExecSuite
@@ -2544,7 +2584,7 @@ class AdaptiveQueryExecSuite
runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " +
"JOIN skewData3 ON value1 = value3")
val shuffles2 = collect(adaptive2) {
Expand All @@ -1954,7 +1944,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
}
assert(shuffles2.size == 4)
val smj2 = findTopLevelSortMergeJoin(adaptive2)
@@ -2703,7 +2744,8 @@ class AdaptiveQueryExecSuite
@@ -2703,7 +2743,8 @@ class AdaptiveQueryExecSuite
}
}

Expand Down
Loading