address review comments

apache · May 16, 2024 · f5380e5 · f5380e5
1 parent e844a0e
commit f5380e5
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 3 deletions.
diff --git a/spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometSparkSessionExtensions.scala b/spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometSparkSessionExtensions.scala
@@ -64,4 +64,6 @@ object ShimCometSparkSessionExtensions {
     }
     true
   }
+
+  protected val EXTENDED_EXPLAIN_PROVIDERS_KEY = "spark.sql.extendedExplainProviders"
 }
diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
@@ -1399,7 +1399,7 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
       CometConf.COMET_EXEC_ENABLED.key -> "true",
       CometConf.COMET_SHUFFLE_ENFORCE_MODE_ENABLED.key -> "true",
       CometConf.COMET_EXEC_ALL_OPERATOR_ENABLED.key -> "true",
-      "spark.sql.extendedExplainProviders" -> "org.apache.comet.ExtendedExplainInfo") {
+      EXTENDED_EXPLAIN_PROVIDERS_KEY -> "org.apache.comet.ExtendedExplainInfo") {
       val table = "test"
       withTable(table) {
         sql(s"create table $table(c0 int, c1 int , c2 float) using parquet")

diff --git a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala
@@ -249,7 +249,7 @@ abstract class CometTestBase
     var dfSpark: Dataset[Row] = null
     withSQLConf(
       CometConf.COMET_ENABLED.key -> "false",
-      "spark.sql.extendedExplainProviders" -> "") {
+      EXTENDED_EXPLAIN_PROVIDERS_KEY -> "") {
       dfSpark = Dataset.ofRows(spark, df.logicalPlan)
       expected = dfSpark.collect()
     }
@@ -259,7 +259,7 @@ abstract class CometTestBase
       dfSpark.queryExecution.explainString(ExtendedMode),
       dfComet.queryExecution.explainString(ExtendedMode))
     if (supportsExtendedExplainInfo(dfSpark.queryExecution)) {
-      assert(expectedInfo.exists(s => diff.contains(s)))
+      assert(expectedInfo.forall(s => diff.contains(s)))
     }
     val extendedInfo =
       new ExtendedExplainInfo().generateExtendedInfo(dfComet.queryExecution.executedPlan)
-Original file line number
+Diff line change
@@ Expand Up / @@ -64,4 +64,6 @@ object ShimCometSparkSessionExtensions { @@
         }
         true
       }
+      protected val EXTENDED_EXPLAIN_PROVIDERS_KEY = "spark.sql.extendedExplainProviders"
     }