chore: Enable Comet explicitly in CometTPCDSQueryTestSuite (#1559)

Enable Comet support explicitly rather than rely on default setting.

Author: andygrove

spark/src/test/scala/org/apache/spark/sql/CometTPCDSQueryTestSuite.scala

@@ -29,6 +29,8 @@ import org.apache.spark.sql.catalyst.util.{fileToString, resourceToString, strin
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.TestSparkSession
 
+import org.apache.comet.CometConf
+
 /**
  * Because we need to modify some methods of Spark `TPCDSQueryTestSuite` but they are private, we
  * copy Spark `TPCDSQueryTestSuite`.
@@ -164,12 +166,18 @@ class CometTPCDSQueryTestSuite extends QueryTest with TPCDSBase with CometSQLQue
     }
   }
 
+  val baseConf: Map[String, String] = Map(
+    CometConf.COMET_ENABLED.key -> "true",
+    CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "true",
+    CometConf.COMET_EXEC_ENABLED.key -> "true",
+    CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true")
+
   val sortMergeJoinConf: Map[String, String] = Map(
     SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
     SQLConf.PREFER_SORTMERGEJOIN.key -> "true")
 
-  val broadcastHashJoinConf: Map[String, String] = Map(
-    SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10485760")
+  val broadcastHashJoinConf: Map[String, String] =
+    Map(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10485760")
 
   val shuffledHashJoinConf: Map[String, String] = Map(
     SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
@@ -213,7 +221,7 @@ class CometTPCDSQueryTestSuite extends QueryTest with TPCDSBase with CometSQLQue
           // that can cause OOM in GitHub Actions
           if (!(sortMergeJoin && name == "q72")) {
             System.gc() // Workaround for GitHub Actions memory limitation, see also SPARK-37368
-            runQuery(queryString, goldenFile, conf)
+            runQuery(queryString, goldenFile, baseConf ++ conf)
           }
         }
       }
@@ -231,7 +239,7 @@ class CometTPCDSQueryTestSuite extends QueryTest with TPCDSBase with CometSQLQue
           // that can cause OOM in GitHub Actions
           if (!(sortMergeJoin && name == "q72")) {
             System.gc() // SPARK-37368
-            runQuery(queryString, goldenFile, conf)
+            runQuery(queryString, goldenFile, baseConf ++ conf)
           }
         }
       }
@@ -245,7 +253,7 @@ class CometTPCDSQueryTestSuite extends QueryTest with TPCDSBase with CometSQLQue
         val goldenFile = new File(s"$baseResourcePath/extended", s"$name.sql.out")
         joinConfs.foreach { conf =>
           System.gc() // SPARK-37368
-          runQuery(queryString, goldenFile, conf)
+          runQuery(queryString, goldenFile, baseConf ++ conf)
         }
       }
     }

spark/src/test/scala/org/apache/spark/sql/CometTPCHQuerySuite.scala

@@ -90,10 +90,8 @@ class CometTPCHQuerySuite extends QueryTest with TPCBase with ShimCometTPCHQuery
     conf.set(CometConf.COMET_EXEC_ENABLED.key, "true")
     conf.set(CometConf.COMET_NATIVE_SCAN_ENABLED.key, "true")
     conf.set(CometConf.COMET_EXEC_SHUFFLE_ENABLED.key, "true")
-    conf.set(CometConf.COMET_SHUFFLE_MODE.key, "jvm")
     conf.set(MEMORY_OFFHEAP_ENABLED.key, "true")
     conf.set(MEMORY_OFFHEAP_SIZE.key, "2g")
-    conf.set(CometConf.COMET_MEMORY_OVERHEAD.key, "2g")
   }
 
   protected override def createSparkSession: TestSparkSession = {

spark/src/test/scala/org/apache/spark/sql/comet/ParquetEncryptionITCase.scala

@@ -30,7 +30,7 @@ import org.scalatest.Tag
 import org.scalatestplus.junit.JUnitRunner
 
 import org.apache.spark.{DebugFilesystem, SparkConf}
-import org.apache.spark.sql.{QueryTest, SparkSession, SQLContext}
+import org.apache.spark.sql.{CometTestBase, SparkSession, SQLContext}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SQLTestUtils
 
@@ -41,7 +41,7 @@ import org.apache.comet.{CometConf, CometSparkSessionExtensions, IntegrationTest
  */
 @RunWith(classOf[JUnitRunner])
 @IntegrationTestSuite
-class ParquetEncryptionITCase extends QueryTest with SQLTestUtils {
+class ParquetEncryptionITCase extends CometTestBase with SQLTestUtils {
   private val encoder = Base64.getEncoder
   private val footerKey =
     encoder.encodeToString("0123456789012345".getBytes(StandardCharsets.UTF_8))
@@ -81,7 +81,12 @@ class ParquetEncryptionITCase extends QueryTest with SQLTestUtils {
             val parquetDF = spark.read.parquet(parquetDir)
             assert(parquetDF.inputFiles.nonEmpty)
             val readDataset = parquetDF.select("a", "b", "c")
-            checkAnswer(readDataset, inputDF)
+
+            if (CometConf.COMET_ENABLED.get(conf)) {
+              checkSparkAnswerAndOperator(readDataset)
+            } else {
+              checkAnswer(readDataset, inputDF)
+            }
           }
         }
     }
@@ -118,19 +123,24 @@ class ParquetEncryptionITCase extends QueryTest with SQLTestUtils {
             val parquetDF = spark.read.parquet(parquetDir)
             assert(parquetDF.inputFiles.nonEmpty)
             val readDataset = parquetDF.select("a", "b", "c")
-            checkAnswer(readDataset, inputDF)
+
+            if (CometConf.COMET_ENABLED.get(conf)) {
+              checkSparkAnswerAndOperator(readDataset)
+            } else {
+              checkAnswer(readDataset, inputDF)
+            }
           }
         }
     }
   }
 
-  protected def sparkConf: SparkConf = {
+  protected override def sparkConf: SparkConf = {
     val conf = new SparkConf()
     conf.set("spark.hadoop.fs.file.impl", classOf[DebugFilesystem].getName)
     conf
   }
 
-  protected def createSparkSession: SparkSession = {
+  protected override def createSparkSession: SparkSession = {
     SparkSession
       .builder()
       .config(sparkConf)
@@ -159,8 +169,8 @@ class ParquetEncryptionITCase extends QueryTest with SQLTestUtils {
   }
 
   private var _spark: SparkSession = _
-  protected implicit def spark: SparkSession = _spark
-  protected implicit def sqlContext: SQLContext = _spark.sqlContext
+  protected implicit override def spark: SparkSession = _spark
+  protected implicit override def sqlContext: SQLContext = _spark.sqlContext
 
   /**
    * Verify that the directory contains an encrypted parquet in encrypted footer mode by means of