test: run Comet in Spark SQL and Iceberg suites that build their own sessions

dev/diffs/3.4.3.diff

@@ -37,6 +37,48 @@ index d3544881af1..0c0618a5e80 100644
      </dependencies>
    </dependencyManagement>
 
+diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala
+index 1cbd9a61289..74aa1f684c9 100644
+--- a/project/SparkBuild.scala
++++ b/project/SparkBuild.scala
+@@ -445,9 +445,11 @@ object SparkBuild extends PomBuild {
+
+   /* Spark SQL Core console settings */
+   enable(SQL.settings)(sql)
++  enable(CometTestSettings.settings)(sql)
+
+   /* Hive console settings */
+   enable(Hive.settings)(hive)
++  enable(CometTestSettings.settings)(hive)
+
+   enable(SparkConnectCommon.settings)(connectCommon)
+   enable(SparkConnect.settings)(connect)
+@@ -1183,6 +1185,25 @@ object SQL {
+   )
+ }
+
++/**
++ * Comet disables itself when its shuffle manager is not registered. When Comet is enabled, make
++ * that shuffle manager the default for every test SparkConf in the modules that have Comet on
++ * their classpath, so that suites which build their own SparkSession or SparkContext run Comet
++ * too, not just the ones that go through SharedSparkSession or TestHive.
++ */
++object CometTestSettings {
++  private val cometEnabled =
++    sys.env.get("ENABLE_COMET").forall(v => v == "1" || v.equalsIgnoreCase("true"))
++
++  lazy val settings: Seq[Setting[_]] =
++    if (cometEnabled && !sys.props.contains("spark.shuffle.manager")) {
++      Seq((Test / javaOptions) +=
++        "-Dspark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
++    } else {
++      Seq.empty
++    }
++}
++
+ object Hive {
+
+   lazy val settings = Seq(
 diff --git a/sql/core/pom.xml b/sql/core/pom.xml
 index b386d135da1..46449e3f3f1 100644
 --- a/sql/core/pom.xml
@@ -1056,6 +1098,50 @@ index 2dabcf01be7..8fcec0d1ce4 100644
        }
      }
    }
+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..a164e273b76 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
+@@ -165,7 +165,15 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper {
+     }
+   }
+
++  // Comet: Comet's own columnar rules take over these plans, so the rules the tests below
++  // inject never see the Spark operators they act on or check for.
++  private def assumeCometDisabled(): Unit = {
++    assume(!org.apache.spark.sql.SparkSession.isCometEnabled,
++      "Skipped when Comet is enabled: Comet replaces the operators the injected rules act on")
++  }
++
+   test("inject adaptive query prep rule") {
++    assumeCometDisabled()
+     val extensions = create { extensions =>
+       // inject rule that will run during AQE query stage preparation and will add custom tags
+       // to the plan
+@@ -259,6 +267,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper {
+   }
+
+   private def testInjectColumnar(enableAQE: Boolean): Unit = {
++    assumeCometDisabled()
+     def collectPlanSteps(plan: SparkPlan): Seq[Int] = plan match {
+       case a: AdaptiveSparkPlanExec =>
+         assert(a.toString.startsWith("AdaptiveSparkPlan isFinalPlan=true"))
+@@ -314,6 +323,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper {
+   }
+
+   test("reset column vectors") {
++    assumeCometDisabled()
+     val session = SparkSession.builder()
+       .master("local[1]")
+       .config(COLUMN_BATCH_SIZE.key, 2)
+@@ -482,6 +492,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper {
+   }
+
+   test("SPARK-38697: Extend SparkSessionExtensions to inject rules into AQE Optimizer") {
++    assumeCometDisabled()
+     def executedPlan(df: Dataset[java.lang.Long]): SparkPlan = {
+       assert(df.queryExecution.executedPlan.isInstanceOf[AdaptiveSparkPlanExec])
+       df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec].executedPlan
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala
 index 18123a4d6ec..0fe185baa33 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala
@@ -1328,6 +1414,84 @@ index c0ec8a58bd5..4e8bc6ed3c5 100644
 
      // Fail to read ancient datetime values.
      withSQLConf(SQLConf.PARQUET_REBASE_MODE_IN_READ.key -> EXCEPTION.toString) {
+diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala
+index 24a98dd83f3..4db28a7a1ce 100644
+--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala
++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala
+@@ -51,6 +51,11 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite {
+     }
+   }
+
++  // Comet: AQE coalesces post-shuffle partitions by map output size, and Comet's shuffle
++  // writes Arrow IPC blocks whose sizes differ from Spark's, so the expected partition
++  // counts below only hold for Spark's shuffle.
++  private def isCometEnabled: Boolean = org.apache.spark.sql.SparkSession.isCometEnabled
++
+   val numInputPartitions: Int = 10
+
+   def withSparkSession(
+@@ -163,9 +168,11 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite {
+             assert(shuffleReads.isEmpty)
+
+           case None =>
+-            assert(shuffleReads.length === 2)
+-            shuffleReads.foreach { read =>
+-              assert(read.outputPartitioning.numPartitions === 2)
++            if (!isCometEnabled) {
++              assert(shuffleReads.length === 2)
++              shuffleReads.foreach { read =>
++                assert(read.outputPartitioning.numPartitions === 2)
++              }
+             }
+         }
+       }
+@@ -214,9 +221,11 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite {
+             assert(shuffleReads.isEmpty)
+
+           case None =>
+-            assert(shuffleReads.length === 2)
+-            shuffleReads.foreach { read =>
+-              assert(read.outputPartitioning.numPartitions === 2)
++            if (!isCometEnabled) {
++              assert(shuffleReads.length === 2)
++              shuffleReads.foreach { read =>
++                assert(read.outputPartitioning.numPartitions === 2)
++              }
+             }
+         }
+       }
+@@ -265,9 +274,11 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite {
+             assert(shuffleReads.isEmpty)
+
+           case None =>
+-            assert(shuffleReads.length === 2)
+-            shuffleReads.foreach { read =>
+-              assert(read.outputPartitioning.numPartitions === 3)
++            if (!isCometEnabled) {
++              assert(shuffleReads.length === 2)
++              shuffleReads.foreach { read =>
++                assert(read.outputPartitioning.numPartitions === 3)
++              }
+             }
+         }
+       }
+@@ -412,10 +423,12 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite {
+       // aggregate on the other side of the union.
+       val finalPlan = resultDf.queryExecution.executedPlan
+         .asInstanceOf[AdaptiveSparkPlanExec].executedPlan
+-      assert(
+-        finalPlan.collect {
+-          case r @ CoalescedShuffleRead() => r
+-        }.size == 2)
++      if (!isCometEnabled) {
++        assert(
++          finalPlan.collect {
++            case r @ CoalescedShuffleRead() => r
++          }.size == 2)
++      }
+     }
+     withSparkSession(test, 100, None)
+   }
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala
 index 418ca3430bb..eb8267192f8 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala
@@ -2433,6 +2597,29 @@ index 3a0bd35cb70..b28f06a757f 100644
      withTempPath { workDir =>
        val workDirPath = workDir.getAbsolutePath
        val input = spark.range(5).toDF("id")
+diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala
+index 6333808b420..81b2704300c 100644
+--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala
++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala
+@@ -21,7 +21,7 @@ import scala.reflect.ClassTag
+
+ import org.apache.spark.AccumulatorSuite
+ import org.apache.spark.internal.config.EXECUTOR_MEMORY
+-import org.apache.spark.sql.{Dataset, QueryTest, Row, SparkSession}
++import org.apache.spark.sql.{Dataset, IgnoreComet, QueryTest, Row, SparkSession}
+ import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BitwiseAnd, BitwiseOr, Cast, Expression, Literal, ShiftLeft}
+ import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide}
+ import org.apache.spark.sql.catalyst.plans.Inner
+@@ -486,7 +486,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils
+     }
+   }
+
+-  test("broadcast join where streamed side's output partitioning is PartitioningCollection") {
++  test("broadcast join where streamed side's output partitioning is PartitioningCollection",
++    IgnoreComet("Comet replaces the join and shuffle operators this test inspects")) {
+     withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "500") {
+       val t1 = (0 until 100).map(i => (i % 5, i % 13)).toDF("i1", "j1")
+       val t2 = (0 until 100).map(i => (i % 5, i % 14)).toDF("i2", "j2")
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
 index 26e61c6b58d..cb09d7e116a 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala