fix: do not run Comet in on-heap mode without spark.comet.exec.onHeap.enabled

.github/workflows/pr_build_linux.yml

@@ -543,6 +543,8 @@ jobs:
               org.apache.spark.CometPluginsMemoryOverheadWarningSuite
               org.apache.spark.CometPluginsMemoryPoolFractionWarningSuite
               org.apache.spark.CometPluginsUnifiedModeSuite
+              org.apache.spark.CometPluginsExtensionOnlySuite
+              org.apache.spark.CometPluginsSparkShuffleManagerSuite
               org.apache.comet.rules.CometScanRuleSuite
               org.apache.comet.rules.CometScanContribSuite
               org.apache.comet.rules.CometScanSchemeFallbackSuite

.github/workflows/pr_build_macos.yml

@@ -191,6 +191,8 @@ jobs:
               org.apache.spark.CometPluginsMemoryOverheadWarningSuite
               org.apache.spark.CometPluginsMemoryPoolFractionWarningSuite
               org.apache.spark.CometPluginsUnifiedModeSuite
+              org.apache.spark.CometPluginsExtensionOnlySuite
+              org.apache.spark.CometPluginsSparkShuffleManagerSuite
               org.apache.comet.rules.CometScanRuleSuite
               org.apache.comet.rules.CometScanContribSuite
               org.apache.comet.rules.CometScanSchemeFallbackSuite

benchmarks/README.md

@@ -71,7 +71,9 @@ $SPARK_HOME/bin/spark-submit \
     --driver-memory 8G \
     --conf spark.driver.memory=8G \
     --conf spark.executor.instances=1 \
-    --conf spark.executor.memory=32G \
+    --conf spark.executor.memory=16G \
+    --conf spark.memory.offHeap.enabled=true \
+    --conf spark.memory.offHeap.size=16G \
     --conf spark.executor.cores=8 \
     --conf spark.cores.max=8 \
     --conf spark.task.cpus=1 \

docs/source/contributor-guide/spark-sql-tests.md

@@ -105,7 +105,8 @@ ENABLE_COMET=true ENABLE_COMET_ONHEAP=true sbt -J-Xmx4096m -Dspark.test.includeS
 -Dspark.comet.enabled=true -Dspark.comet.debug.enabled=true -Dspark.plugins=org.apache.spark.CometPlugin -DXmx4096m -Dspark.executor.heartbeatInterval=20000 -Dspark.network.timeout=10000 --add-exports=java.base/sun.nio.ch=ALL-UNNAMED --add-opens=java.base/java.nio=ALL-UNNAMED
 ```
 
-2. Set `ENABLE_COMET=true` in environment variables
+2. Set `ENABLE_COMET=true` and `ENABLE_COMET_ONHEAP=true` in environment variables. Without
+   `ENABLE_COMET_ONHEAP`, Comet stays disabled, because the tests do not enable off-heap memory.
    ![img.png](img.png)
 3. After the above tests are configured, spark tests can be run with debugging enabled on spark/comet code. Note that Comet is added as a dependency and the classes are readonly while debugging from Spark. Any new changes to Comet are to be built and deployed locally through the command (`PROFILES="-Pspark-3.4" make release`)
 

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

@@ -130,7 +130,18 @@ object CometSparkSessionExtensions extends Logging {
       return false
     }
 
-    if (COMET_SHUFFLE_ENABLED.get(conf) && !isCometShuffleManagerEnabled(conf)) {
+    // CometDriverPlugin makes the same check before registering this extension, but an
+    // application can also register the extension directly with spark.sql.extensions. The memory
+    // mode comes from the SparkContext's conf, which is what executors use. A session's SQLConf
+    // can disagree: when the SparkContext already exists, SparkSession.Builder copies core
+    // configs into it without applying them.
+    val offHeapEnabled = Option(SparkEnv.get).exists(env => isOffHeapEnabled(env.conf))
+    if (!offHeapEnabled && !COMET_ONHEAP_ENABLED.get(conf)) {
+      logWarning("Comet extension is disabled because Spark is not running in off-heap mode.")
+      return false
+    }
+
+    if (COMET_SHUFFLE_ENABLED.get(conf) && !isCometShuffleManagerEnabled) {
       logWarning(
         "Comet extension is disabled because spark.shuffle.manager is not set to " +
           s"${classOf[CometShuffleManager].getName} or " +
@@ -174,7 +185,7 @@ object CometSparkSessionExtensions extends Logging {
   // dependencies without passing through ordinary exchange selection. Celeborn requires explicit
   // native opt-in and compatible application settings; local Comet shuffle keeps its behavior.
   def isCometShuffleEnabled(conf: SQLConf): Boolean =
-    COMET_SHUFFLE_ENABLED.get(conf) && isCometShuffleManagerEnabled(conf) &&
+    COMET_SHUFFLE_ENABLED.get(conf) && isCometShuffleManagerEnabled &&
       cometCelebornShuffleFallbackReason(conf, numPartitions = 1).isEmpty
 
   private def activeCelebornShuffleManager: Option[CometCelebornShuffleManager] =
@@ -187,14 +198,14 @@ object CometSparkSessionExtensions extends Logging {
       conf.getConfString(SHUFFLE_MANAGER_KEY, "") ==
       classOf[CometCelebornShuffleManager].getName
 
-  def isCometShuffleManagerEnabled(conf: SQLConf): Boolean = {
-    conf.contains(SHUFFLE_MANAGER_KEY) && {
-      val manager = conf.getConfString(SHUFFLE_MANAGER_KEY)
-      manager == classOf[CometShuffleManager].getName ||
-      manager == classOf[CometCelebornShuffleManager].getName ||
-      activeCelebornShuffleManager.isDefined
+  // Inspect the manager SparkEnv holds rather than spark.shuffle.manager in the session's
+  // SQLConf. The manager is created once for the application, and when the SparkContext already
+  // exists, SparkSession.Builder copies core configs into the SQLConf without applying them.
+  def isCometShuffleManagerEnabled: Boolean =
+    Option(SparkEnv.get).flatMap(env => Option(env.shuffleManager)).exists {
+      case _: CometShuffleManager | _: CometCelebornShuffleManager => true
+      case _ => false
     }
-  }
 
   /**
    * Native mode and execution can be chosen per query. Encryption, stage recovery, and Celeborn

spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala

@@ -843,7 +843,7 @@ object CometShuffleExchangeExec
   private def isCometShuffleEnabledReason(op: SparkPlan): Option[String] = {
     if (!COMET_SHUFFLE_ENABLED.get(op.conf)) {
       Some(s"Comet shuffle is not enabled: ${COMET_SHUFFLE_ENABLED.key} is not enabled")
-    } else if (!isCometShuffleManagerEnabled(op.conf)) {
+    } else if (!isCometShuffleManagerEnabled) {
       Some(
         s"spark.shuffle.manager is not set to ${classOf[CometShuffleManager].getName} or " +
           classOf[CometCelebornShuffleManager].getName)

spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala

@@ -62,23 +62,19 @@ class CometSparkSessionExtensionsSuite extends CometTestBase {
     NativeBase.setLoaded(true)
   }
 
-  test("isCometLoaded requires CometShuffleManager when shuffle.enabled=true") {
-    val conf = new SQLConf
-    conf.setConfString(CometConf.COMET_ENABLED.key, "true")
-
-    // Default: shuffle.enabled=true. Without spark.shuffle.manager set, Comet must be disabled.
-    assert(!isCometLoaded(conf))
-
-    // Opt out: shuffle.enabled=false. Comet should load (assumes native lib is available).
-    conf.setConfString(CometConf.COMET_SHUFFLE_ENABLED.key, "false")
-    assert(isCometLoaded(conf))
-
-    // shuffle.enabled=true with the Comet shuffle manager registered: Comet should load.
-    conf.setConfString(CometConf.COMET_SHUFFLE_ENABLED.key, "true")
-    conf.setConfString(
-      "spark.shuffle.manager",
-      "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
-    assert(isCometLoaded(conf))
+  test("isCometLoaded follows the application's shuffle manager, not the session conf") {
+    // This suite's SparkContext runs CometShuffleManager. A session conf naming another manager,
+    // which SparkSession.Builder can leave behind, does not change what runs the shuffle.
+    Seq(
+      "org.apache.spark.shuffle.sort.SortShuffleManager",
+      "org.apache.spark.shuffle.celeborn.SparkShuffleManager").foreach { manager =>
+      val conf = new SQLConf
+      conf.setConfString(CometConf.COMET_ENABLED.key, "true")
+      conf.setConfString(CometConf.COMET_SHUFFLE_ENABLED.key, "true")
+      conf.setConfString("spark.shuffle.manager", manager)
+      assert(isCometShuffleEnabled(conf), manager)
+      assert(isCometLoaded(conf), manager)
+    }
   }
 
   test("the composite manager is recognized without requiring the optional Celeborn client") {
@@ -87,24 +83,12 @@ class CometSparkSessionExtensionsSuite extends CometTestBase {
     conf.setConfString(CometConf.COMET_SHUFFLE_ENABLED.key, "true")
     conf.setConfString(CometConf.COMET_SHUFFLE_MODE.key, "native")
     conf.setConfString("spark.shuffle.manager", classOf[CometCelebornShuffleManager].getName)
-    assert(isCometShuffleManagerEnabled(conf))
+    assert(isCometShuffleManagerEnabled)
     assert(isCometLoaded(conf))
     // A session-only setting must not replace this suite's actual local shuffle manager.
     assert(!isCometShuffleEnabled(conf))
   }
 
-  test("the stock Celeborn manager cannot accept Comet shuffle dependencies") {
-    val conf = new SQLConf
-    conf.setConfString(CometConf.COMET_ENABLED.key, "true")
-    conf.setConfString(CometConf.COMET_SHUFFLE_ENABLED.key, "true")
-    conf.setConfString(
-      "spark.shuffle.manager",
-      "org.apache.spark.shuffle.celeborn.SparkShuffleManager")
-    assert(!isCometShuffleManagerEnabled(conf))
-    assert(!isCometShuffleEnabled(conf))
-    assert(!isCometLoaded(conf))
-  }
-
   test("local auto mode retains Comet columnar fallback for unsupported native partitioning") {
     withSQLConf(
       CometConf.COMET_SHUFFLE_MODE.key -> "auto",

spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala

@@ -22,10 +22,12 @@ package org.apache.spark
 import java.io.File
 
 import org.apache.logging.log4j.Level
-import org.apache.spark.sql.{CometTestBase, SaveMode}
+import org.apache.spark.sql.{CometTestBase, SaveMode, SparkSession}
+import org.apache.spark.sql.comet.CometPlan
+import org.apache.spark.sql.comet.execution.shuffle.{CometShuffleExchangeExec, CometShuffleManager}
 import org.apache.spark.sql.internal.StaticSQLConf
 
-import org.apache.comet.COMET_VERSION
+import org.apache.comet.{COMET_VERSION, CometConf}
 
 class CometPluginsSuite extends CometTestBase {
   override protected def sparkConf: SparkConf = {
@@ -253,3 +255,105 @@ class CometPluginsUnifiedModeSuite extends CometTestBase {
     assert(execMemOverhead4 == "1G")
   }
 }
+
+class CometPluginsExtensionOnlySuite extends CometTestBase {
+  // No plugin and no off-heap memory. CometTestBase registers CometSparkSessionExtensions
+  // directly, as an application can with spark.sql.extensions.
+  override protected def sparkConf: SparkConf = {
+    val conf = new SparkConf()
+    conf.set(
+      "spark.shuffle.manager",
+      "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
+    conf.set("spark.comet.enabled", "true")
+    conf.set("spark.comet.exec.enabled", "true")
+    // Set explicitly, since ENABLE_COMET_ONHEAP in the environment changes the default.
+    conf.set(CometConf.COMET_ONHEAP_ENABLED.key, "false")
+    conf
+  }
+
+  private val query = "SELECT _1 FROM tbl WHERE _1 > 5"
+
+  test("Comet is disabled when off-heap memory is disabled") {
+    // Listen on the package logger, as CometExecRuleSuite does. For a logger with no config of its
+    // own, withLogAppender creates one that outlives the test and does not pass events up, so
+    // listening on CometSparkSessionExtensions directly would hide its warnings from later
+    // appenders on org.apache.comet.
+    val appender = new LogAppender("off-heap mode warning")
+    withParquetTable((0 until 10).map(i => (i, i.toString)), "tbl") {
+      withLogAppender(appender, Seq("org.apache.comet"), Some(Level.WARN)) {
+        val (_, plan) = checkSparkAnswer(query)
+        assert(collect(plan) { case op: CometPlan => op }.isEmpty, plan)
+      }
+    }
+    assert(
+      appender.loggingEvents
+        .exists(_.getMessage.getFormattedMessage.contains("not running in off-heap mode")))
+  }
+
+  test("Comet stays disabled when only the session conf enables off-heap memory") {
+    withParquetTable((0 until 10).map(i => (i, i.toString)), "tbl") {
+      // The session already exists, so the builder only copies the setting into its SQLConf.
+      // The SparkContext, whose conf the executors use, stays on-heap.
+      val session =
+        SparkSession.builder().config("spark.memory.offHeap.enabled", "true").getOrCreate()
+      try {
+        assert(session eq spark)
+        assert(spark.sessionState.conf.getConfString("spark.memory.offHeap.enabled") == "true")
+        val (_, plan) = checkSparkAnswer(query)
+        assert(collect(plan) { case op: CometPlan => op }.isEmpty, plan)
+      } finally {
+        spark.sessionState.conf.unsetConf("spark.memory.offHeap.enabled")
+      }
+    }
+  }
+
+  test("spark.comet.exec.onHeap.enabled enables Comet without off-heap memory") {
+    withParquetTable((0 until 10).map(i => (i, i.toString)), "tbl") {
+      withSQLConf(CometConf.COMET_ONHEAP_ENABLED.key -> "true") {
+        val (_, plan) = checkSparkAnswer(query)
+        assert(collect(plan) { case op: CometPlan => op }.nonEmpty, plan)
+      }
+    }
+  }
+}
+
+class CometPluginsSparkShuffleManagerSuite extends CometTestBase {
+  // The application runs Spark's own shuffle manager.
+  override protected def sparkConf: SparkConf = {
+    val conf = new SparkConf()
+    conf.set("spark.memory.offHeap.enabled", "true")
+    conf.set("spark.memory.offHeap.size", "2g")
+    conf.set("spark.comet.enabled", "true")
+    conf.set("spark.comet.exec.enabled", "true")
+    conf
+  }
+
+  private val query = "SELECT _2, count(*) FROM tbl GROUP BY _2"
+
+  test("Comet stays disabled when only the session conf names the Comet shuffle manager") {
+    withParquetTable((0 until 100).map(i => (i, (i % 7).toString)), "tbl") {
+      // The session already exists, so the builder only copies the setting into its SQLConf.
+      // Spark's shuffle manager still runs the shuffle, and it cannot read a Comet shuffle.
+      val manager = classOf[CometShuffleManager].getName
+      val session = SparkSession.builder().config("spark.shuffle.manager", manager).getOrCreate()
+      try {
+        assert(session eq spark)
+        assert(spark.sessionState.conf.getConfString("spark.shuffle.manager") == manager)
+        val (_, plan) = checkSparkAnswer(query)
+        assert(collect(plan) { case op: CometPlan => op }.isEmpty, plan)
+      } finally {
+        spark.sessionState.conf.unsetConf("spark.shuffle.manager")
+      }
+    }
+  }
+
+  test("Comet runs with Spark's shuffle when Comet shuffle is disabled") {
+    withParquetTable((0 until 100).map(i => (i, (i % 7).toString)), "tbl") {
+      withSQLConf(CometConf.COMET_SHUFFLE_ENABLED.key -> "false") {
+        val (_, plan) = checkSparkAnswer(query)
+        assert(collect(plan) { case op: CometPlan => op }.nonEmpty, plan)
+        assert(collect(plan) { case op: CometShuffleExchangeExec => op }.isEmpty, plan)
+      }
+    }
+  }
+}

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

@@ -51,6 +51,9 @@ trait CometTPCQueryBase extends Logging {
       .set(
         "spark.shuffle.manager",
         "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
+      // Off-heap memory is disabled, so Comet runs in on-heap mode, which must be enabled
+      // explicitly or Comet stays disabled.
+      .set("spark.comet.exec.onHeap.enabled", "true")
 
     val sparkSession = SparkSession
       .builder()

spark/src/test/scala/org/apache/spark/sql/benchmark/CometBenchmarkBase.scala

@@ -62,6 +62,9 @@ trait CometBenchmarkBase
       .set(
         "spark.shuffle.manager",
         "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
+      // Off-heap memory is disabled, so Comet runs in on-heap mode, which must be enabled
+      // explicitly or Comet stays disabled.
+      .set("spark.comet.exec.onHeap.enabled", "true")
 
     val sparkSession = SparkSession
       .builder()

spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastHashJoinBenchmark.scala

@@ -43,6 +43,7 @@ object CometBroadcastHashJoinBenchmark extends CometBenchmarkBase {
       .set(
         "spark.shuffle.manager",
         "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
+      .set("spark.comet.exec.onHeap.enabled", "true")
 
     val sparkSession = SparkSession
       .builder()

spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastNestedLoopJoinBenchmark.scala

@@ -44,6 +44,7 @@ object CometBroadcastNestedLoopJoinBenchmark extends CometBenchmarkBase {
       .set(
         "spark.shuffle.manager",
         "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
+      .set("spark.comet.exec.onHeap.enabled", "true")
 
     val sparkSession = SparkSession
       .builder()

spark/src/test/scala/org/apache/spark/sql/benchmark/CometExecBenchmark.scala

@@ -46,6 +46,7 @@ object CometExecBenchmark extends CometBenchmarkBase {
       .set(
         "spark.shuffle.manager",
         "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
+      .set("spark.comet.exec.onHeap.enabled", "true")
       .set("spark.comet.shuffle.jvm.spillThreshold", "30000")
 
     val sparkSession = SparkSession

spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashJoinBenchmark.scala

@@ -43,6 +43,7 @@ object CometHashJoinBenchmark extends CometBenchmarkBase {
       .set(
         "spark.shuffle.manager",
         "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
+      .set("spark.comet.exec.onHeap.enabled", "true")
 
     val sparkSession = SparkSession
       .builder()

spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleBenchmark.scala

@@ -89,6 +89,7 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
       .set(
         "spark.shuffle.manager",
         "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
+      .set("spark.comet.exec.onHeap.enabled", "true")
       .set("spark.comet.shuffle.jvm.spillThreshold", "30000")
 
     val sparkSession = SparkSession

spark/src/test/scala/org/apache/spark/sql/benchmark/CometSortMergeJoinBenchmark.scala

@@ -43,6 +43,7 @@ object CometSortMergeJoinBenchmark extends CometBenchmarkBase {
       .set(
         "spark.shuffle.manager",
         "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
+      .set("spark.comet.exec.onHeap.enabled", "true")
 
     val sparkSession = SparkSession
       .builder()

spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala

@@ -144,7 +144,7 @@ class CometCelebornShufflePlanningSuite extends CometTestBase {
 
   test("the actual composite manager loads Comet and default auto preserves Spark shuffle") {
     val conf = spark.sessionState.conf
-    assert(isCometShuffleManagerEnabled(conf))
+    assert(isCometShuffleManagerEnabled)
     assertNativeExecutionLoaded()
     assert(CometConf.COMET_SHUFFLE_MODE.get(conf) == "auto")
     assert(!isCometShuffleEnabled(conf))