chore: Enable Comet shuffle with AQE coalesce partitions

.github/workflows/spark_sql_test.yml

@@ -43,7 +43,7 @@ jobs:
   spark-sql-catalyst:
     strategy:
       matrix:
-        os: [ubuntu-latest]
+        os: [ubuntu-24.04]
         java-version: [11]
         spark-version: [{short: '3.4', full: '3.4.3'}, {short: '3.5', full: '3.5.1'}]
         module:
@@ -76,7 +76,7 @@ jobs:
         run: |
           cd apache-spark
           rm -rf /root/.m2/repository/org/apache/parquet # somehow parquet cache requires cleanups
-          ENABLE_COMET=true build/sbt ${{ matrix.module.args1 }} "${{ matrix.module.args2 }}"
+          ENABLE_COMET=true ENABLE_COMET_SHUFFLE=${{ matrix.module.name == 'sql/core-1' && 'false' || 'true' }} build/sbt ${{ matrix.module.args1 }} "${{ matrix.module.args2 }}"
         env:
           LC_ALL: "C.UTF-8"
 

.github/workflows/spark_sql_test_ansi.yml

@@ -41,7 +41,7 @@ jobs:
   spark-sql-catalyst:
     strategy:
       matrix:
-        os: [ubuntu-latest]
+        os: [ubuntu-24.04]
         java-version: [17]
         spark-version: [{short: '4.0', full: '4.0.0-preview1'}]
         module:
@@ -74,7 +74,7 @@ jobs:
         run: |
           cd apache-spark
           rm -rf /root/.m2/repository/org/apache/parquet # somehow parquet cache requires cleanups
-          RUST_BACKTRACE=1 ENABLE_COMET=true ENABLE_COMET_ANSI_MODE=true build/sbt ${{ matrix.module.args1 }} "${{ matrix.module.args2 }}"
+          RUST_BACKTRACE=1 ENABLE_COMET=true ENABLE_COMET_ANSI_MODE=true ENABLE_COMET_SHUFFLE=${{ matrix.module.name == 'sql/core-1' && 'false' || 'true' }} build/sbt ${{ matrix.module.args1 }} "${{ matrix.module.args2 }}"
         env:
           LC_ALL: "C.UTF-8"
 

common/src/main/scala/org/apache/comet/CometConf.scala

@@ -204,17 +204,6 @@ object CometConf extends ShimCometConf {
     .checkValues(Set("native", "jvm", "auto"))
     .createWithDefault("jvm")
 
-  val COMET_SHUFFLE_ENFORCE_MODE_ENABLED: ConfigEntry[Boolean] =
-    conf("spark.comet.shuffle.enforceMode.enabled")
-      .doc(
-        "Comet shuffle doesn't support Spark AQE coalesce partitions. If AQE coalesce " +
-          "partitions is enabled, Comet shuffle won't be triggered even enabled. This config " +
-          "is used to enforce Comet to trigger shuffle even if AQE coalesce partitions is " +
-          "enabled. This is for testing purpose only.")
-      .internal()
-      .booleanConf
-      .createWithDefault(false)
-
   val COMET_EXEC_BROADCAST_FORCE_ENABLED: ConfigEntry[Boolean] =
     conf(s"$COMET_EXEC_CONFIG_PREFIX.broadcast.enabled")
       .doc(

docs/source/contributor-guide/benchmarking.md

@@ -67,7 +67,6 @@ $SPARK_HOME/bin/spark-submit \
     --conf spark.comet.cast.allowIncompatible=true \
     --conf spark.comet.exec.shuffle.enabled=true \
     --conf spark.comet.exec.shuffle.mode=auto \
-    --conf spark.comet.shuffle.enforceMode.enabled=true \
     --conf spark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager \
     tpcbench.py \
     --benchmark tpch \

spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala

@@ -1069,21 +1069,13 @@ object CometSparkSessionExtensions extends Logging {
   }
 
   private[comet] def isCometShuffleEnabled(conf: SQLConf): Boolean =
-    COMET_EXEC_SHUFFLE_ENABLED.get(conf) && isCometShuffleManagerEnabled(conf) &&
-      // TODO: AQE coalesce partitions feature causes Comet shuffle memory leak.
-      // We should disable Comet shuffle when AQE coalesce partitions is enabled.
-      (!conf.coalesceShufflePartitionsEnabled || COMET_SHUFFLE_ENFORCE_MODE_ENABLED.get())
+    COMET_EXEC_SHUFFLE_ENABLED.get(conf) && isCometShuffleManagerEnabled(conf)
 
   private[comet] def getCometShuffleNotEnabledReason(conf: SQLConf): Option[String] = {
     if (!COMET_EXEC_SHUFFLE_ENABLED.get(conf)) {
       Some(s"${COMET_EXEC_SHUFFLE_ENABLED.key} is not enabled")
     } else if (!isCometShuffleManagerEnabled(conf)) {
       Some(s"spark.shuffle.manager is not set to ${CometShuffleManager.getClass.getName}")
-    } else if (conf.coalesceShufflePartitionsEnabled && !COMET_SHUFFLE_ENFORCE_MODE_ENABLED
-        .get()) {
-      Some(
-        s"${SQLConf.COALESCE_PARTITIONS_ENABLED.key} is enabled and " +
-          s"${COMET_SHUFFLE_ENFORCE_MODE_ENABLED.key} is not enabled")
     } else {
       None
     }