fix: Disable Comet shuffle with AQE coalesce partitions enabled

apache · May 4, 2024 · 130f8b7 · 130f8b7
1 parent 26af020
commit 130f8b7
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 1 deletion.
diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
@@ -916,7 +916,9 @@ object CometSparkSessionExtensions extends Logging {
   private[comet] def isCometShuffleEnabled(conf: SQLConf): Boolean =
     COMET_EXEC_SHUFFLE_ENABLED.get(conf) &&
       (conf.contains("spark.shuffle.manager") && conf.getConfString("spark.shuffle.manager") ==
-        "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
+        "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") &&
+      // TODO: AQE coalesce partitions feature causes Comet columnar shuffle memory leak
+      !conf.coalesceShufflePartitionsEnabled
 
   private[comet] def isCometScanEnabled(conf: SQLConf): Boolean = {
     COMET_SCAN_ENABLED.get(conf)

diff --git a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala
@@ -64,6 +64,21 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar
 
   import testImplicits._
 
+  setupTestData()
+
+  test("Disable Comet shuffle with AQE coalesce partitions enabled") {
+    withSQLConf(
+      CometConf.COMET_EXEC_ENABLED.key -> "true",
+      SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+      SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
+      SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
+      val df = sql(
+        "SELECT * FROM (SELECT * FROM testData WHERE key = 0) t1 FULL JOIN " +
+          "testData2 t2 ON t1.key = t2.a")
+      checkShuffleAnswer(df, 0)
+    }
+  }
+
   test("columnar shuffle on nested struct including nulls") {
     Seq(10, 201).foreach { numPartitions =>
       Seq("1.0", "10.0").foreach { ratio =>

diff --git a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala
@@ -76,6 +76,7 @@ abstract class CometTestBase
     conf.set(MEMORY_OFFHEAP_SIZE.key, "2g")
     conf.set(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key, "1g")
     conf.set(SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key, "1g")
+    conf.set(SQLConf.COALESCE_PARTITIONS_ENABLED.key, "false")
     conf.set(CometConf.COMET_ENABLED.key, "true")
     conf.set(CometConf.COMET_EXEC_ENABLED.key, "true")
     conf.set(CometConf.COMET_EXEC_ALL_OPERATOR_ENABLED.key, "true")