Skip to content

Commit 83ca07a

Browse files
andygroveandygrove
committed
test: run Spark's SPARK-35332 AQE cache test with Comet again
The diffs skipped CachedTableSuite's 'SPARK-35332: Make cache plan disable configs configurable - check AQE' on Spark 4.0 and later, because SparkPlanInfo drew the cached plan only below Spark's own InMemoryTableScanExec. Since #6577 it is drawn below CometInMemoryTableScanExec too, so the test finds the cached plan. The test's check that the cached plan was coalesced, which the diffs already adapt, looked for AQEShuffleRead only below a columnar-to-row transition, where it is when Comet's plan is cached in Spark's format. In Comet's columnar format the cached plan ends in the read itself, as in Spark, so the check accepts it there too. The test passes on Spark 4.0, 4.1 and 4.2 with either format, and with Comet off.
1 parent 52c8df5 commit 83ca07a

3 files changed

Lines changed: 21 additions & 48 deletions

File tree

‎dev/diffs/4.0.4.diff‎

Lines changed: 7 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -375,7 +375,7 @@ index 21a3ce1e122..f4762ab98f0 100644
375375

376376
-- In COMPENSATION views get invalidated if the type can't cast
377377
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
378-
index 0f42502f1d9..81990d4a97f 100644
378+
index 0f42502f1d9..f7ff02bb0df 100644
379379
--- a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
380380
+++ b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
381381
@@ -36,10 +36,11 @@ import org.apache.spark.sql.catalyst.analysis.TempTableAlreadyExistsException
@@ -432,35 +432,26 @@ index 0f42502f1d9..81990d4a97f 100644
432432
val tableName = "ntzCache"
433433
withTable(tableName) {
434434
sql(s"CACHE TABLE $tableName AS SELECT TIMESTAMP_NTZ'2021-01-01 00:00:00'")
435-
@@ -1626,7 +1634,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils
436-
}
437-
}
438-
439-
- test("SPARK-35332: Make cache plan disable configs configurable - check AQE") {
440-
+ test("SPARK-35332: Make cache plan disable configs configurable - check AQE",
441-
+ IgnoreComet("Spark's SQL UI shows a cached plan only under Spark's own cache scan")) {
442-
withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2",
443-
SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1",
444-
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") {
445-
@@ -1659,9 +1668,18 @@ class CachedTableSuite extends QueryTest with SQLTestUtils
435+
@@ -1659,9 +1667,19 @@ class CachedTableSuite extends QueryTest with SQLTestUtils
446436
_.nodeName.contains("TableCacheQueryStage"))
447437
val aqeNode = findNodeInSparkPlanInfo(inMemoryScanNode.get,
448438
_.nodeName.contains("AdaptiveSparkPlan"))
449439
- val aqePlanRoot = findNodeInSparkPlanInfo(inMemoryScanNode.get,
450440
- _.nodeName.contains("ResultQueryStage"))
451441
- aqePlanRoot.get.children.head.nodeName == "AQEShuffleRead"
452442
+ // Spark 4.0 wraps results in ResultQueryStage. The coalescing indicator is AQEShuffleRead
453-
+ // as the direct child of InputAdapter.
443+
+ // as its direct child, as in Spark, when the table is cached in Comet's columnar format,
444+
+ // or as the direct child of InputAdapter when Comet's plan is cached in Spark's format.
454445
+ // AdaptiveSparkPlan -> ResultQueryStage -> WholestageCodegen ->
455446
+ // CometColumnarToRow -> InputAdapter -> AQEShuffleRead (if coalesced)
456447
+ val resultStage = aqeNode.get.children.head // ResultQueryStage
457-
+ val wsc = resultStage.children.head // WholeStageCodegen
448+
+ val wsc = resultStage.children.head // WholeStageCodegen or AQEShuffleRead
458449
+ val c2r = wsc.children.head // ColumnarToRow or CometColumnarToRow
459450
+ val inputAdapter = c2r.children.head // InputAdapter
460-
+ resultStage.nodeName == "ResultQueryStage" &&
451+
+ resultStage.nodeName == "ResultQueryStage" && (wsc.nodeName == "AQEShuffleRead" ||
461452
+ wsc.nodeName.startsWith("WholeStageCodegen") && // could be "WholeStageCodegen (1)"
462453
+ (c2r.nodeName == "CometColumnarToRow" || c2r.nodeName == "ColumnarToRow") &&
463-
+ inputAdapter.children.head.nodeName == "AQEShuffleRead"
454+
+ inputAdapter.children.head.nodeName == "AQEShuffleRead")
464455
}
465456

466457
withTempView("t0", "t1", "t2") {

‎dev/diffs/4.1.3.diff‎

Lines changed: 7 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -423,7 +423,7 @@ index 26d8f750f6e..c888f8e0844 100644
423423

424424
test("SPARK-51777: sql.columnar.* classes registered in KryoSerializer") {
425425
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
426-
index 0d807aeae4d..ce07fba149d 100644
426+
index 0d807aeae4d..bfa1316f471 100644
427427
--- a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
428428
+++ b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
429429
@@ -37,6 +37,7 @@ import org.apache.spark.sql.catalyst.analysis.TempTableAlreadyExistsException
@@ -484,35 +484,26 @@ index 0d807aeae4d..ce07fba149d 100644
484484
val tableName = "ntzCache"
485485
withTable(tableName) {
486486
sql(s"CACHE TABLE $tableName AS SELECT TIMESTAMP_NTZ'2021-01-01 00:00:00'")
487-
@@ -1640,7 +1648,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils
488-
}
489-
}
490-
491-
- test("SPARK-35332: Make cache plan disable configs configurable - check AQE") {
492-
+ test("SPARK-35332: Make cache plan disable configs configurable - check AQE",
493-
+ IgnoreComet("Spark's SQL UI shows a cached plan only under Spark's own cache scan")) {
494-
withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2",
495-
SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1",
496-
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") {
497-
@@ -1673,9 +1682,18 @@ class CachedTableSuite extends QueryTest with SQLTestUtils
487+
@@ -1673,9 +1681,19 @@ class CachedTableSuite extends QueryTest with SQLTestUtils
498488
_.nodeName.contains("TableCacheQueryStage"))
499489
val aqeNode = findNodeInSparkPlanInfo(inMemoryScanNode.get,
500490
_.nodeName.contains("AdaptiveSparkPlan"))
501491
- val aqePlanRoot = findNodeInSparkPlanInfo(inMemoryScanNode.get,
502492
- _.nodeName.contains("ResultQueryStage"))
503493
- aqePlanRoot.get.children.head.nodeName == "AQEShuffleRead"
504494
+ // Spark 4.0 wraps results in ResultQueryStage. The coalescing indicator is AQEShuffleRead
505-
+ // as the direct child of InputAdapter.
495+
+ // as its direct child, as in Spark, when the table is cached in Comet's columnar format,
496+
+ // or as the direct child of InputAdapter when Comet's plan is cached in Spark's format.
506497
+ // AdaptiveSparkPlan -> ResultQueryStage -> WholestageCodegen ->
507498
+ // CometColumnarToRow -> InputAdapter -> AQEShuffleRead (if coalesced)
508499
+ val resultStage = aqeNode.get.children.head // ResultQueryStage
509-
+ val wsc = resultStage.children.head // WholeStageCodegen
500+
+ val wsc = resultStage.children.head // WholeStageCodegen or AQEShuffleRead
510501
+ val c2r = wsc.children.head // ColumnarToRow or CometColumnarToRow
511502
+ val inputAdapter = c2r.children.head // InputAdapter
512-
+ resultStage.nodeName == "ResultQueryStage" &&
503+
+ resultStage.nodeName == "ResultQueryStage" && (wsc.nodeName == "AQEShuffleRead" ||
513504
+ wsc.nodeName.startsWith("WholeStageCodegen") && // could be "WholeStageCodegen (1)"
514505
+ (c2r.nodeName == "CometColumnarToRow" || c2r.nodeName == "ColumnarToRow") &&
515-
+ inputAdapter.children.head.nodeName == "AQEShuffleRead"
506+
+ inputAdapter.children.head.nodeName == "AQEShuffleRead")
516507
}
517508

518509
withTempView("t0", "t1", "t2") {

‎dev/diffs/4.2.0.diff‎

Lines changed: 7 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -441,7 +441,7 @@ index 72a2da16054..54146036f8a 100644
441441

442442
test("SPARK-51777: sql.columnar.* classes registered in KryoSerializer") {
443443
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
444-
index 085dbcd8046..ac98c64ae92 100644
444+
index 085dbcd8046..8cee9832317 100644
445445
--- a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
446446
+++ b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
447447
@@ -37,6 +37,7 @@ import org.apache.spark.sql.catalyst.analysis.TempTableAlreadyExistsException
@@ -502,35 +502,26 @@ index 085dbcd8046..ac98c64ae92 100644
502502
val tableName = "ntzCache"
503503
withTable(tableName) {
504504
sql(s"CACHE TABLE $tableName AS SELECT TIMESTAMP_NTZ'2021-01-01 00:00:00'")
505-
@@ -1670,7 +1678,8 @@ class CachedTableSuite extends SharedSparkSession
506-
}
507-
}
508-
509-
- test("SPARK-35332: Make cache plan disable configs configurable - check AQE") {
510-
+ test("SPARK-35332: Make cache plan disable configs configurable - check AQE",
511-
+ IgnoreComet("Spark's SQL UI shows a cached plan only under Spark's own cache scan")) {
512-
withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2",
513-
SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1",
514-
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") {
515-
@@ -1703,9 +1712,18 @@ class CachedTableSuite extends SharedSparkSession
505+
@@ -1703,9 +1711,19 @@ class CachedTableSuite extends SharedSparkSession
516506
_.nodeName.contains("TableCacheQueryStage"))
517507
val aqeNode = findNodeInSparkPlanInfo(inMemoryScanNode.get,
518508
_.nodeName.contains("AdaptiveSparkPlan"))
519509
- val aqePlanRoot = findNodeInSparkPlanInfo(inMemoryScanNode.get,
520510
- _.nodeName.contains("ResultQueryStage"))
521511
- aqePlanRoot.get.children.head.nodeName == "AQEShuffleRead"
522512
+ // Spark 4.0 wraps results in ResultQueryStage. The coalescing indicator is AQEShuffleRead
523-
+ // as the direct child of InputAdapter.
513+
+ // as its direct child, as in Spark, when the table is cached in Comet's columnar format,
514+
+ // or as the direct child of InputAdapter when Comet's plan is cached in Spark's format.
524515
+ // AdaptiveSparkPlan -> ResultQueryStage -> WholestageCodegen ->
525516
+ // CometColumnarToRow -> InputAdapter -> AQEShuffleRead (if coalesced)
526517
+ val resultStage = aqeNode.get.children.head // ResultQueryStage
527-
+ val wsc = resultStage.children.head // WholeStageCodegen
518+
+ val wsc = resultStage.children.head // WholeStageCodegen or AQEShuffleRead
528519
+ val c2r = wsc.children.head // ColumnarToRow or CometColumnarToRow
529520
+ val inputAdapter = c2r.children.head // InputAdapter
530-
+ resultStage.nodeName == "ResultQueryStage" &&
521+
+ resultStage.nodeName == "ResultQueryStage" && (wsc.nodeName == "AQEShuffleRead" ||
531522
+ wsc.nodeName.startsWith("WholeStageCodegen") && // could be "WholeStageCodegen (1)"
532523
+ (c2r.nodeName == "CometColumnarToRow" || c2r.nodeName == "ColumnarToRow") &&
533-
+ inputAdapter.children.head.nodeName == "AQEShuffleRead"
524+
+ inputAdapter.children.head.nodeName == "AQEShuffleRead")
534525
}
535526

536527
withTempView("t0", "t1", "t2") {

0 commit comments

Comments
 (0)