From f5380e556ce08d7d887cbffca0c6376de474d99d Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Thu, 16 May 2024 10:08:27 -0700
Subject: [PATCH] address review comments

---
 .../apache/comet/shims/ShimCometSparkSessionExtensions.scala  | 2 ++
 .../test/scala/org/apache/comet/CometExpressionSuite.scala    | 2 +-
 spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala | 4 ++--
 3 files changed, 5 insertions(+), 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
index ffec1bd40..22c567bb4 100644
--- 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
index 40d15281b..f3fd50e9e 100644
--- 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
index 31c53508c..68905ac34 100644
--- 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)