From f68737949f3c535ce8d493623495b96fe45fb884 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Tue, 24 Dec 2024 04:08:16 -0800
Subject: [PATCH 1/2]  test: Enable Comet by default and ignore some tests in
 SparkSessionExtensionSuite

---
 dev/diffs/3.4.3.diff          | 79 +++++++++++++++++++++++++------
 dev/diffs/3.5.1.diff          | 89 +++++++++++++++++++++++++++--------
 dev/diffs/4.0.0-preview1.diff | 75 ++++++++++++++++++++++++-----
 3 files changed, 195 insertions(+), 48 deletions(-)

diff --git a/dev/diffs/3.4.3.diff b/dev/diffs/3.4.3.diff
index 12f739848..8fe012d62 100644
--- a/dev/diffs/3.4.3.diff
+++ b/dev/diffs/3.4.3.diff
@@ -1,5 +1,5 @@
 diff --git a/pom.xml b/pom.xml
-index d3544881af1..bf0e2b53c70 100644
+index d3544881af1..26ab186c65d 100644
 --- a/pom.xml
 +++ b/pom.xml
 @@ -148,6 +148,8 @@
@@ -38,7 +38,7 @@ index d3544881af1..bf0e2b53c70 100644
    </dependencyManagement>
  
 diff --git a/sql/core/pom.xml b/sql/core/pom.xml
-index b386d135da1..854aec17c2d 100644
+index b386d135da1..46449e3f3f1 100644
 --- a/sql/core/pom.xml
 +++ b/sql/core/pom.xml
 @@ -77,6 +77,10 @@
@@ -53,7 +53,7 @@ index b386d135da1..854aec17c2d 100644
      <!--
        This spark-tags test-dep is needed even though it isn't used in this module, otherwise testing-cmds that exclude
 diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala b/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala
-index c595b50950b..6b60213e775 100644
+index c595b50950b..3abb6cb9441 100644
 --- a/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala
 +++ b/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala
 @@ -102,7 +102,7 @@ class SparkSession private(
@@ -79,7 +79,7 @@ index c595b50950b..6b60213e775 100644
    }
  
 +  private def loadCometExtension(sparkContext: SparkContext): Seq[String] = {
-+    if (sparkContext.getConf.getBoolean("spark.comet.enabled", false)) {
++    if (sparkContext.getConf.getBoolean("spark.comet.enabled", isCometEnabled)) {
 +      Seq("org.apache.comet.CometSparkSessionExtensions")
 +    } else {
 +      Seq.empty
@@ -100,6 +100,19 @@ index c595b50950b..6b60213e775 100644
        try {
          val extensionConfClass = Utils.classForName(extensionConfClassName)
          val extensionConf = extensionConfClass.getConstructor().newInstance()
+@@ -1323,4 +1333,12 @@ object SparkSession extends Logging {
+       }
+     }
+   }
++
++  /**
++   * Whether Comet extension is enabled
++   */
++  def isCometEnabled: Boolean = {
++    val v = System.getenv("ENABLE_COMET")
++    v == null || v.toBoolean
++  }
+ }
 diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala
 index db587dd9868..aac7295a53d 100644
 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala
@@ -957,6 +970,46 @@ index 525d97e4998..8a3e7457618 100644
      AccumulatorSuite.verifyPeakExecutionMemorySet(sparkContext, "external sort") {
        sql("SELECT * FROM testData2 ORDER BY a ASC, b ASC").collect()
      }
+diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
+index 48ad10992c5..3ec91cccfa8 100644
+--- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
++++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
+@@ -188,11 +188,13 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper {
+     }
+   }
+ 
+-  test("inject columnar AQE on") {
++  test("inject columnar AQE on",
++    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
+     testInjectColumnar(true)
+   }
+ 
+-  test("inject columnar AQE off") {
++  test("inject columnar AQE off",
++    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
+     testInjectColumnar(false)
+   }
+ 
+@@ -213,7 +215,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper {
+     }
+   }
+ 
+-  test("SPARK-39991: AQE should retain column statistics from completed query stages") {
++  test("SPARK-39991: AQE should retain column statistics from completed query stages",
++    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
+     val extensions = create { extensions =>
+       extensions.injectColumnar(_ =>
+         MyColumnarRule(PreRuleReplaceAddWithBrokenVersion(), MyPostRule()))
+@@ -313,7 +316,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper {
+     }
+   }
+ 
+-  test("reset column vectors") {
++  test("reset column vectors",
++    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
+     val session = SparkSession.builder()
+       .master("local[1]")
+       .config(COLUMN_BATCH_SIZE.key, 2)
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala
 index 75eabcb96f2..36e3318ad7e 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala
@@ -2720,7 +2773,7 @@ index abe606ad9c1..2d930b64cca 100644
      val tblTargetName = "tbl_target"
      val tblSourceQualified = s"default.$tblSourceName"
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala
-index dd55fcfe42c..aa9b0be8e68 100644
+index dd55fcfe42c..2702f87c1f1 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala
 @@ -41,6 +41,7 @@ import org.apache.spark.sql.catalyst.plans.PlanTest
@@ -2744,17 +2797,14 @@ index dd55fcfe42c..aa9b0be8e68 100644
      }
    }
  
-@@ -242,6 +247,41 @@ private[sql] trait SQLTestUtilsBase
+@@ -242,6 +247,38 @@ private[sql] trait SQLTestUtilsBase
      protected override def _sqlContext: SQLContext = self.spark.sqlContext
    }
  
 +  /**
 +   * Whether Comet extension is enabled
 +   */
-+  protected def isCometEnabled: Boolean = {
-+    val v = System.getenv("ENABLE_COMET")
-+    v != null && v.toBoolean
-+  }
++  protected def isCometEnabled: Boolean = SparkSession.isCometEnabled
 +
 +  /**
 +   * Whether to enable ansi mode This is only effective when
@@ -2786,7 +2836,7 @@ index dd55fcfe42c..aa9b0be8e68 100644
    protected override def withSQLConf(pairs: (String, String)*)(f: => Unit): Unit = {
      SparkSession.setActiveSession(spark)
      super.withSQLConf(pairs: _*)(f)
-@@ -434,6 +474,8 @@ private[sql] trait SQLTestUtilsBase
+@@ -434,6 +471,8 @@ private[sql] trait SQLTestUtilsBase
      val schema = df.schema
      val withoutFilters = df.queryExecution.executedPlan.transform {
        case FilterExec(_, child) => child
@@ -2884,10 +2934,10 @@ index 1966e1e64fd..cde97a0aafe 100644
        spark.sql(
          """
 diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala
-index 07361cfdce9..6673c141c9a 100644
+index 07361cfdce9..e40c59a4207 100644
 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala
 +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala
-@@ -55,25 +55,53 @@ object TestHive
+@@ -55,25 +55,52 @@ object TestHive
      new SparkContext(
        System.getProperty("spark.sql.test.master", "local[1]"),
        "TestSQLContext",
@@ -2929,8 +2979,7 @@ index 07361cfdce9..6673c141c9a 100644
 +          // ConstantPropagation etc.
 +          .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName)
 +
-+        val v = System.getenv("ENABLE_COMET")
-+        if (v != null && v.toBoolean) {
++        if (SparkSession.isCometEnabled) {
 +          conf
 +            .set("spark.sql.extensions", "org.apache.comet.CometSparkSessionExtensions")
 +            .set("spark.comet.enabled", "true")
diff --git a/dev/diffs/3.5.1.diff b/dev/diffs/3.5.1.diff
index 64d6681bc..9d123975c 100644
--- a/dev/diffs/3.5.1.diff
+++ b/dev/diffs/3.5.1.diff
@@ -1,5 +1,5 @@
 diff --git a/pom.xml b/pom.xml
-index 0f504dbee85..f6019da888a 100644
+index 0f504dbee85..430ec217e59 100644
 --- a/pom.xml
 +++ b/pom.xml
 @@ -152,6 +152,8 @@
@@ -38,7 +38,7 @@ index 0f504dbee85..f6019da888a 100644
    </dependencyManagement>
  
 diff --git a/sql/core/pom.xml b/sql/core/pom.xml
-index c46ab7b8fce..d8b99c2c115 100644
+index c46ab7b8fce..13357e8c7a6 100644
 --- a/sql/core/pom.xml
 +++ b/sql/core/pom.xml
 @@ -77,6 +77,10 @@
@@ -53,7 +53,7 @@ index c46ab7b8fce..d8b99c2c115 100644
      <!--
        This spark-tags test-dep is needed even though it isn't used in this module, otherwise testing-cmds that exclude
 diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala b/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala
-index 27ae10b3d59..064cbc252ea 100644
+index 27ae10b3d59..78e69902dfd 100644
 --- a/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala
 +++ b/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala
 @@ -1353,6 +1353,14 @@ object SparkSession extends Logging {
@@ -61,11 +61,11 @@ index 27ae10b3d59..064cbc252ea 100644
    }
  
 +  private def loadCometExtension(sparkContext: SparkContext): Seq[String] = {
-+    if (sparkContext.getConf.getBoolean("spark.comet.enabled", false)) {
-+        Seq("org.apache.comet.CometSparkSessionExtensions")
-+      } else {
-+        Seq.empty
-+      }
++    if (sparkContext.getConf.getBoolean("spark.comet.enabled", isCometEnabled)) {
++      Seq("org.apache.comet.CometSparkSessionExtensions")
++    } else {
++      Seq.empty
++    }
 +  }
 +
    /**
@@ -79,6 +79,19 @@ index 27ae10b3d59..064cbc252ea 100644
      extensionConfClassNames.foreach { extensionConfClassName =>
        try {
          val extensionConfClass = Utils.classForName(extensionConfClassName)
+@@ -1396,4 +1405,12 @@ object SparkSession extends Logging {
+       }
+     }
+   }
++
++  /**
++   * Whether Comet extension is enabled
++   */
++  def isCometEnabled: Boolean = {
++    val v = System.getenv("ENABLE_COMET")
++    v == null || v.toBoolean
++  }
+ }
 diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala
 index db587dd9868..aac7295a53d 100644
 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala
@@ -959,6 +972,46 @@ index cfeccbdf648..803d8734cc4 100644
      AccumulatorSuite.verifyPeakExecutionMemorySet(sparkContext, "external sort") {
        sql("SELECT * FROM testData2 ORDER BY a ASC, b ASC").collect()
      }
+diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
+index 8b4ac474f87..4d92f325889 100644
+--- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
++++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
+@@ -190,11 +190,13 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt
+     }
+   }
+ 
+-  test("inject columnar AQE on") {
++  test("inject columnar AQE on",
++    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
+     testInjectColumnar(true)
+   }
+ 
+-  test("inject columnar AQE off") {
++  test("inject columnar AQE off",
++    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
+     testInjectColumnar(false)
+   }
+ 
+@@ -215,7 +217,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt
+     }
+   }
+ 
+-  test("SPARK-39991: AQE should retain column statistics from completed query stages") {
++  test("SPARK-39991: AQE should retain column statistics from completed query stages",
++    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
+     val extensions = create { extensions =>
+       extensions.injectColumnar(_ =>
+         MyColumnarRule(PreRuleReplaceAddWithBrokenVersion(), MyPostRule()))
+@@ -315,7 +318,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt
+     }
+   }
+ 
+-  test("reset column vectors") {
++  test("reset column vectors",
++    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
+     val session = SparkSession.builder()
+       .master("local[1]")
+       .config(COLUMN_BATCH_SIZE.key, 2)
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala
 index fbc256b3396..0821999c7c2 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala
@@ -1310,7 +1363,7 @@ index 47679ed7865..9ffbaecb98e 100644
      assert(collectWithSubqueries(plan) { case s: SortAggregateExec => s }.length == sortAggCount)
    }
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala
-index 5a413c77754..c52f4b3818c 100644
+index 5a413c77754..a6f97dccb67 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala
 @@ -20,6 +20,7 @@ package org.apache.spark.sql.execution
@@ -2705,7 +2758,7 @@ index abe606ad9c1..2d930b64cca 100644
      val tblTargetName = "tbl_target"
      val tblSourceQualified = s"default.$tblSourceName"
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala
-index dd55fcfe42c..aa9b0be8e68 100644
+index dd55fcfe42c..2702f87c1f1 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala
 @@ -41,6 +41,7 @@ import org.apache.spark.sql.catalyst.plans.PlanTest
@@ -2729,17 +2782,14 @@ index dd55fcfe42c..aa9b0be8e68 100644
      }
    }
  
-@@ -242,6 +247,41 @@ private[sql] trait SQLTestUtilsBase
+@@ -242,6 +247,38 @@ private[sql] trait SQLTestUtilsBase
      protected override def _sqlContext: SQLContext = self.spark.sqlContext
    }
  
 +  /**
 +   * Whether Comet extension is enabled
 +   */
-+  protected def isCometEnabled: Boolean = {
-+    val v = System.getenv("ENABLE_COMET")
-+    v != null && v.toBoolean
-+  }
++  protected def isCometEnabled: Boolean = SparkSession.isCometEnabled
 +
 +  /**
 +   * Whether to enable ansi mode This is only effective when
@@ -2771,7 +2821,7 @@ index dd55fcfe42c..aa9b0be8e68 100644
    protected override def withSQLConf(pairs: (String, String)*)(f: => Unit): Unit = {
      SparkSession.setActiveSession(spark)
      super.withSQLConf(pairs: _*)(f)
-@@ -434,6 +474,8 @@ private[sql] trait SQLTestUtilsBase
+@@ -434,6 +471,8 @@ private[sql] trait SQLTestUtilsBase
      val schema = df.schema
      val withoutFilters = df.queryExecution.executedPlan.transform {
        case FilterExec(_, child) => child
@@ -2869,10 +2919,10 @@ index dc8b184fcee..dd69a989d40 100644
        spark.sql(
          """
 diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala
-index 9284b35fb3e..e8984be5ebc 100644
+index 9284b35fb3e..2a0269bdc16 100644
 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala
 +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala
-@@ -53,25 +53,53 @@ object TestHive
+@@ -53,25 +53,52 @@ object TestHive
      new SparkContext(
        System.getProperty("spark.sql.test.master", "local[1]"),
        "TestSQLContext",
@@ -2914,8 +2964,7 @@ index 9284b35fb3e..e8984be5ebc 100644
 +          // ConstantPropagation etc.
 +          .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName)
 +
-+        val v = System.getenv("ENABLE_COMET")
-+        if (v != null && v.toBoolean) {
++        if (SparkSession.isCometEnabled) {
 +          conf
 +            .set("spark.sql.extensions", "org.apache.comet.CometSparkSessionExtensions")
 +            .set("spark.comet.enabled", "true")
diff --git a/dev/diffs/4.0.0-preview1.diff b/dev/diffs/4.0.0-preview1.diff
index db62ed607..b3744c2eb 100644
--- a/dev/diffs/4.0.0-preview1.diff
+++ b/dev/diffs/4.0.0-preview1.diff
@@ -53,7 +53,7 @@ index 19f6303be36..6c0e77882e6 100644
      <!--
        This spark-tags test-dep is needed even though it isn't used in this module, otherwise testing-cmds that exclude
 diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala b/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala
-index 466e4cf8131..7ead2715b5d 100644
+index 466e4cf8131..798f118464c 100644
 --- a/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala
 +++ b/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala
 @@ -1376,6 +1376,14 @@ object SparkSession extends Logging {
@@ -61,7 +61,7 @@ index 466e4cf8131..7ead2715b5d 100644
    }
  
 +  private def loadCometExtension(sparkContext: SparkContext): Seq[String] = {
-+    if (sparkContext.getConf.getBoolean("spark.comet.enabled", false)) {
++    if (sparkContext.getConf.getBoolean("spark.comet.enabled", isCometEnabled)) {
 +      Seq("org.apache.comet.CometSparkSessionExtensions")
 +    } else {
 +      Seq.empty
@@ -81,6 +81,19 @@ index 466e4cf8131..7ead2715b5d 100644
        try {
          val extensionConfClass = Utils.classForName(extensionConfClassName)
          val extensionConf = extensionConfClass.getConstructor().newInstance()
+@@ -1420,4 +1429,12 @@ object SparkSession extends Logging {
+       }
+     }
+   }
++
++  /**
++   * Whether Comet extension is enabled
++   */
++  def isCometEnabled: Boolean = {
++    val v = System.getenv("ENABLE_COMET")
++    v == null || v.toBoolean
++  }
+ }
 diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala
 index 7c45b02ee84..9f2b608c9f5 100644
 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala
@@ -1031,6 +1044,46 @@ index 56c364e2084..fc3abd7cdc4 100644
      withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false",
        SQLConf.ANSI_ENABLED.key -> "true") {
        withTable("t") {
+diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
+index 4d38e360f43..21e88695976 100644
+--- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
++++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
+@@ -190,11 +190,13 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt
+     }
+   }
+ 
+-  test("inject columnar AQE on") {
++  test("inject columnar AQE on",
++    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
+     testInjectColumnar(true)
+   }
+ 
+-  test("inject columnar AQE off") {
++  test("inject columnar AQE off",
++    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
+     testInjectColumnar(false)
+   }
+ 
+@@ -215,7 +217,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt
+     }
+   }
+ 
+-  test("SPARK-39991: AQE should retain column statistics from completed query stages") {
++  test("SPARK-39991: AQE should retain column statistics from completed query stages",
++    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
+     val extensions = create { extensions =>
+       extensions.injectColumnar(_ =>
+         MyColumnarRule(PreRuleReplaceAddWithBrokenVersion(), MyPostRule()))
+@@ -315,7 +318,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt
+     }
+   }
+ 
+-  test("reset column vectors") {
++  test("reset column vectors",
++    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
+     val session = SparkSession.builder()
+       .master("local[1]")
+       .config(COLUMN_BATCH_SIZE.key, 2)
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala
 index 68f14f13bbd..174636cefb5 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala
@@ -2910,7 +2963,7 @@ index af07aceaed1..ed0b5e6d9be 100644
      val tblTargetName = "tbl_target"
      val tblSourceQualified = s"default.$tblSourceName"
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala
-index 5fbf379644f..51debb8c481 100644
+index 5fbf379644f..1f2e5d7cfa0 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala
 @@ -41,6 +41,7 @@ import org.apache.spark.sql.catalyst.plans.PlanTest
@@ -2934,17 +2987,14 @@ index 5fbf379644f..51debb8c481 100644
      }
    }
  
-@@ -243,6 +248,42 @@ private[sql] trait SQLTestUtilsBase
+@@ -243,6 +248,39 @@ private[sql] trait SQLTestUtilsBase
      protected override def _sqlContext: SQLContext = self.spark.sqlContext
    }
  
 +  /**
 +   * Whether Comet extension is enabled
 +   */
-+  protected def isCometEnabled: Boolean = {
-+    val v = System.getenv("ENABLE_COMET")
-+    v != null && v.toBoolean
-+  }
++  protected def isCometEnabled: Boolean = SparkSession.isCometEnabled
 +
 +  /**
 +   * Whether to enable ansi mode This is only effective when
@@ -2977,7 +3027,7 @@ index 5fbf379644f..51debb8c481 100644
    protected override def withSQLConf[T](pairs: (String, String)*)(f: => T): T = {
      SparkSession.setActiveSession(spark)
      super.withSQLConf(pairs: _*)(f)
-@@ -434,6 +475,8 @@ private[sql] trait SQLTestUtilsBase
+@@ -434,6 +472,8 @@ private[sql] trait SQLTestUtilsBase
      val schema = df.schema
      val withoutFilters = df.queryExecution.executedPlan.transform {
        case FilterExec(_, child) => child
@@ -3061,10 +3111,10 @@ index 52abd248f3a..7a199931a08 100644
          case d: DynamicPruningExpression => d.child
        }
 diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala
-index 3f8de93b330..fc06e81f41a 100644
+index 3f8de93b330..2912f98cc5b 100644
 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala
 +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala
-@@ -53,24 +53,53 @@ object TestHive
+@@ -53,24 +53,52 @@ object TestHive
      new SparkContext(
        System.getProperty("spark.sql.test.master", "local[1]"),
        "TestSQLContext",
@@ -3106,8 +3156,7 @@ index 3f8de93b330..fc06e81f41a 100644
 +          // ConstantPropagation etc.
 +          .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName)
 +
-+        val v = System.getenv("ENABLE_COMET")
-+        if (v != null && v.toBoolean) {
++        if (SparkSession.isCometEnabled) {
 +          conf
 +            .set("spark.sql.extensions", "org.apache.comet.CometSparkSessionExtensions")
 +            .set("spark.comet.enabled", "true")

From 758458aa2d87f766578d95b73ff70e0459cfe174 Mon Sep 17 00:00:00 2001
From: Kazuyuki Tanimura <ktanimura@apple.com>
Date: Fri, 3 Jan 2025 14:40:00 -0800
Subject: [PATCH 2/2]  test: Enable Comet by default and ignore some tests in
 SparkSessionExtensionSuite

---
 dev/diffs/3.4.3.diff          | 59 +++++++++++++++--------------------
 dev/diffs/3.5.1.diff          | 59 +++++++++++++++--------------------
 dev/diffs/4.0.0-preview1.diff | 59 +++++++++++++++--------------------
 3 files changed, 75 insertions(+), 102 deletions(-)

diff --git a/dev/diffs/3.4.3.diff b/dev/diffs/3.4.3.diff
index 8fe012d62..7b4cd8f76 100644
--- a/dev/diffs/3.4.3.diff
+++ b/dev/diffs/3.4.3.diff
@@ -971,45 +971,36 @@ index 525d97e4998..8a3e7457618 100644
        sql("SELECT * FROM testData2 ORDER BY a ASC, b ASC").collect()
      }
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
-index 48ad10992c5..3ec91cccfa8 100644
+index 48ad10992c5..51d1ee65422 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
-@@ -188,11 +188,13 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper {
-     }
-   }
- 
--  test("inject columnar AQE on") {
-+  test("inject columnar AQE on",
-+    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
-     testInjectColumnar(true)
-   }
- 
--  test("inject columnar AQE off") {
-+  test("inject columnar AQE off",
-+    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
-     testInjectColumnar(false)
-   }
- 
-@@ -213,7 +215,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper {
-     }
-   }
- 
--  test("SPARK-39991: AQE should retain column statistics from completed query stages") {
-+  test("SPARK-39991: AQE should retain column statistics from completed query stages",
-+    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
-     val extensions = create { extensions =>
-       extensions.injectColumnar(_ =>
-         MyColumnarRule(PreRuleReplaceAddWithBrokenVersion(), MyPostRule()))
-@@ -313,7 +316,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper {
-     }
-   }
- 
--  test("reset column vectors") {
-+  test("reset column vectors",
-+    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
+@@ -221,6 +221,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper {
+     withSession(extensions) { session =>
+       session.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED, true)
+       session.conf.set(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key, "-1")
++      // https://github.com/apache/datafusion-comet/issues/1197
++      session.conf.set("spark.comet.enabled", false)
+       assert(session.sessionState.columnarRules.contains(
+         MyColumnarRule(PreRuleReplaceAddWithBrokenVersion(), MyPostRule())))
+       import session.sqlContext.implicits._
+@@ -279,6 +281,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper {
+     }
+     withSession(extensions) { session =>
+       session.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED, enableAQE)
++      // https://github.com/apache/datafusion-comet/issues/1197
++      session.conf.set("spark.comet.enabled", false)
+       assert(session.sessionState.columnarRules.contains(
+         MyColumnarRule(PreRuleReplaceAddWithBrokenVersion(), MyPostRule())))
+       import session.sqlContext.implicits._
+@@ -317,6 +321,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper {
      val session = SparkSession.builder()
        .master("local[1]")
        .config(COLUMN_BATCH_SIZE.key, 2)
++      // https://github.com/apache/datafusion-comet/issues/1197
++      .config("spark.comet.enabled", false)
+       .withExtensions { extensions =>
+         extensions.injectColumnar(session =>
+           MyColumnarRule(PreRuleReplaceAddWithBrokenVersion(), MyPostRule())) }
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala
 index 75eabcb96f2..36e3318ad7e 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala
diff --git a/dev/diffs/3.5.1.diff b/dev/diffs/3.5.1.diff
index 9d123975c..1e12ee33c 100644
--- a/dev/diffs/3.5.1.diff
+++ b/dev/diffs/3.5.1.diff
@@ -973,45 +973,36 @@ index cfeccbdf648..803d8734cc4 100644
        sql("SELECT * FROM testData2 ORDER BY a ASC, b ASC").collect()
      }
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
-index 8b4ac474f87..4d92f325889 100644
+index 8b4ac474f87..3f79f20822f 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
-@@ -190,11 +190,13 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt
-     }
-   }
- 
--  test("inject columnar AQE on") {
-+  test("inject columnar AQE on",
-+    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
-     testInjectColumnar(true)
-   }
- 
--  test("inject columnar AQE off") {
-+  test("inject columnar AQE off",
-+    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
-     testInjectColumnar(false)
-   }
- 
-@@ -215,7 +217,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt
-     }
-   }
- 
--  test("SPARK-39991: AQE should retain column statistics from completed query stages") {
-+  test("SPARK-39991: AQE should retain column statistics from completed query stages",
-+    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
-     val extensions = create { extensions =>
-       extensions.injectColumnar(_ =>
-         MyColumnarRule(PreRuleReplaceAddWithBrokenVersion(), MyPostRule()))
-@@ -315,7 +318,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt
-     }
-   }
- 
--  test("reset column vectors") {
-+  test("reset column vectors",
-+    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
+@@ -223,6 +223,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt
+     withSession(extensions) { session =>
+       session.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED, true)
+       session.conf.set(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key, "-1")
++      // https://github.com/apache/datafusion-comet/issues/1197
++      session.conf.set("spark.comet.enabled", false)
+       assert(session.sessionState.columnarRules.contains(
+         MyColumnarRule(PreRuleReplaceAddWithBrokenVersion(), MyPostRule())))
+       import session.sqlContext.implicits._
+@@ -281,6 +283,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt
+     }
+     withSession(extensions) { session =>
+       session.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED, enableAQE)
++      // https://github.com/apache/datafusion-comet/issues/1197
++      session.conf.set("spark.comet.enabled", false)
+       assert(session.sessionState.columnarRules.contains(
+         MyColumnarRule(PreRuleReplaceAddWithBrokenVersion(), MyPostRule())))
+       import session.sqlContext.implicits._
+@@ -319,6 +323,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt
      val session = SparkSession.builder()
        .master("local[1]")
        .config(COLUMN_BATCH_SIZE.key, 2)
++      // https://github.com/apache/datafusion-comet/issues/1197
++      .config("spark.comet.enabled", false)
+       .withExtensions { extensions =>
+         extensions.injectColumnar(session =>
+           MyColumnarRule(PreRuleReplaceAddWithBrokenVersion(), MyPostRule())) }
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala
 index fbc256b3396..0821999c7c2 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala
diff --git a/dev/diffs/4.0.0-preview1.diff b/dev/diffs/4.0.0-preview1.diff
index b3744c2eb..919741ad2 100644
--- a/dev/diffs/4.0.0-preview1.diff
+++ b/dev/diffs/4.0.0-preview1.diff
@@ -1045,45 +1045,36 @@ index 56c364e2084..fc3abd7cdc4 100644
        SQLConf.ANSI_ENABLED.key -> "true") {
        withTable("t") {
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
-index 4d38e360f43..21e88695976 100644
+index 4d38e360f43..3c272af0b62 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
-@@ -190,11 +190,13 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt
-     }
-   }
- 
--  test("inject columnar AQE on") {
-+  test("inject columnar AQE on",
-+    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
-     testInjectColumnar(true)
-   }
- 
--  test("inject columnar AQE off") {
-+  test("inject columnar AQE off",
-+    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
-     testInjectColumnar(false)
-   }
- 
-@@ -215,7 +217,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt
-     }
-   }
- 
--  test("SPARK-39991: AQE should retain column statistics from completed query stages") {
-+  test("SPARK-39991: AQE should retain column statistics from completed query stages",
-+    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
-     val extensions = create { extensions =>
-       extensions.injectColumnar(_ =>
-         MyColumnarRule(PreRuleReplaceAddWithBrokenVersion(), MyPostRule()))
-@@ -315,7 +318,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt
-     }
-   }
- 
--  test("reset column vectors") {
-+  test("reset column vectors",
-+    IgnoreComet("https://github.com/apache/datafusion-comet/issues/1197")) {
+@@ -223,6 +223,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt
+     withSession(extensions) { session =>
+       session.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED, true)
+       session.conf.set(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key, "-1")
++      // https://github.com/apache/datafusion-comet/issues/1197
++      session.conf.set("spark.comet.enabled", false)
+       assert(session.sessionState.columnarRules.contains(
+         MyColumnarRule(PreRuleReplaceAddWithBrokenVersion(), MyPostRule())))
+       import session.sqlContext.implicits._
+@@ -281,6 +283,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt
+     }
+     withSession(extensions) { session =>
+       session.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED, enableAQE)
++      // https://github.com/apache/datafusion-comet/issues/1197
++      session.conf.set("spark.comet.enabled", false)
+       assert(session.sessionState.columnarRules.contains(
+         MyColumnarRule(PreRuleReplaceAddWithBrokenVersion(), MyPostRule())))
+       import session.sqlContext.implicits._
+@@ -319,6 +323,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt
      val session = SparkSession.builder()
        .master("local[1]")
        .config(COLUMN_BATCH_SIZE.key, 2)
++      // https://github.com/apache/datafusion-comet/issues/1197
++      .config("spark.comet.enabled", false)
+       .withExtensions { extensions =>
+         extensions.injectColumnar(session =>
+           MyColumnarRule(PreRuleReplaceAddWithBrokenVersion(), MyPostRule())) }
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala
 index 68f14f13bbd..174636cefb5 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala