From b4e48b7d670640ac8d42af2535dfff377640cb1a Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 24 Sep 2026 09:05:20 -0600 Subject: [PATCH 1/4] fix: don't run Comet in on-heap mode without spark.comet.exec.onHeap.enabled CometDriverPlugin disabled Comet when off-heap memory was disabled and spark.comet.exec.onHeap.enabled was not set, but an application that registered CometSparkSessionExtensions directly with spark.sql.extensions skipped the check and ran Comet in on-heap mode without a warning. isCometLoaded now makes the same check. The benchmarks that register the extension without off-heap memory set spark.comet.exec.onHeap.enabled so that they keep running Comet. --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + benchmarks/README.md | 4 +- .../contributor-guide/spark-sql-tests.md | 3 +- .../comet/CometSparkSessionExtensions.scala | 13 ++++++ .../CometSparkSessionExtensionsSuite.scala | 32 +++++++++++-- .../org/apache/spark/CometPluginsSuite.scala | 45 ++++++++++++++++++- .../apache/spark/sql/CometTPCQueryBase.scala | 3 ++ .../sql/benchmark/CometBenchmarkBase.scala | 3 ++ .../CometBroadcastHashJoinBenchmark.scala | 1 + ...ometBroadcastNestedLoopJoinBenchmark.scala | 1 + .../sql/benchmark/CometExecBenchmark.scala | 1 + .../benchmark/CometHashJoinBenchmark.scala | 1 + .../sql/benchmark/CometShuffleBenchmark.scala | 1 + .../CometSortMergeJoinBenchmark.scala | 1 + 15 files changed, 104 insertions(+), 7 deletions(-) diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 8abd6c92496..fa4167d1c58 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -543,6 +543,7 @@ jobs: org.apache.spark.CometPluginsMemoryOverheadWarningSuite org.apache.spark.CometPluginsMemoryPoolFractionWarningSuite org.apache.spark.CometPluginsUnifiedModeSuite + org.apache.spark.CometPluginsExtensionOnlySuite org.apache.comet.rules.CometScanRuleSuite org.apache.comet.rules.CometScanContribSuite org.apache.comet.rules.CometScanSchemeFallbackSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index af6bbc48c85..b33da808c64 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -191,6 +191,7 @@ jobs: org.apache.spark.CometPluginsMemoryOverheadWarningSuite org.apache.spark.CometPluginsMemoryPoolFractionWarningSuite org.apache.spark.CometPluginsUnifiedModeSuite + org.apache.spark.CometPluginsExtensionOnlySuite org.apache.comet.rules.CometScanRuleSuite org.apache.comet.rules.CometScanContribSuite org.apache.comet.rules.CometScanSchemeFallbackSuite diff --git a/benchmarks/README.md b/benchmarks/README.md index a95d2e2d3d6..0b055900946 100644 --- a/benchmarks/README.md +++ b/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 \ diff --git a/docs/source/contributor-guide/spark-sql-tests.md b/docs/source/contributor-guide/spark-sql-tests.md index 34e09836263..e82683d8b6e 100644 --- a/docs/source/contributor-guide/spark-sql-tests.md +++ b/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`) diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala index 749a3d6ecb0..b48ad248e6a 100644 --- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala +++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala @@ -130,6 +130,13 @@ object CometSparkSessionExtensions extends Logging { return false } + // CometDriverPlugin makes the same check before registering this extension, but an + // application can also register the extension directly with spark.sql.extensions. + if (!isOffHeapEnabled(conf) && !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(conf)) { logWarning( "Comet extension is disabled because spark.shuffle.manager is not set to " + @@ -292,6 +299,12 @@ object CometSparkSessionExtensions extends Logging { sparkConf.getBoolean("spark.memory.offHeap.enabled", false) } + // Spark copies the SparkConf into each session's SQLConf and, by default, refuses to set core + // configs such as this one at session level, so this is the application's memory mode. + private def isOffHeapEnabled(conf: SQLConf): Boolean = { + conf.getConfString("spark.memory.offHeap.enabled", "false").toBoolean + } + /** * Record a fallback reason on a `TreeNode` (a Spark operator or expression) explaining why * Comet cannot accelerate it. Reasons recorded here are surfaced in extended explain output diff --git a/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala b/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala index 331d4b4ece9..334a4383dfd 100644 --- a/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala @@ -36,8 +36,16 @@ class CometSparkSessionExtensionsSuite extends CometTestBase { import CometSparkSessionExtensions._ - test("isCometLoaded") { + // isCometLoaded also requires off-heap memory, so the tests of its other checks start from a + // conf that has it. + private def offHeapSQLConf(): SQLConf = { val conf = new SQLConf + conf.setConfString("spark.memory.offHeap.enabled", "true") + conf + } + + test("isCometLoaded") { + val conf = offHeapSQLConf() // Disable Comet shuffle so this test can focus on other checks without needing // spark.shuffle.manager to be set. conf.setConfString(CometConf.COMET_SHUFFLE_ENABLED.key, "false") @@ -63,7 +71,7 @@ class CometSparkSessionExtensionsSuite extends CometTestBase { } test("isCometLoaded requires CometShuffleManager when shuffle.enabled=true") { - val conf = new SQLConf + val conf = offHeapSQLConf() conf.setConfString(CometConf.COMET_ENABLED.key, "true") // Default: shuffle.enabled=true. Without spark.shuffle.manager set, Comet must be disabled. @@ -81,9 +89,25 @@ class CometSparkSessionExtensionsSuite extends CometTestBase { assert(isCometLoaded(conf)) } - test("the composite manager is recognized without requiring the optional Celeborn client") { + test("isCometLoaded requires off-heap memory unless on-heap mode is enabled") { val conf = new SQLConf conf.setConfString(CometConf.COMET_ENABLED.key, "true") + conf.setConfString(CometConf.COMET_SHUFFLE_ENABLED.key, "false") + // Set explicitly, since ENABLE_COMET_ONHEAP in the environment changes the default. + conf.setConfString(CometConf.COMET_ONHEAP_ENABLED.key, "false") + assert(!isCometLoaded(conf)) + + conf.setConfString("spark.memory.offHeap.enabled", "true") + assert(isCometLoaded(conf)) + + conf.setConfString("spark.memory.offHeap.enabled", "false") + conf.setConfString(CometConf.COMET_ONHEAP_ENABLED.key, "true") + assert(isCometLoaded(conf)) + } + + test("the composite manager is recognized without requiring the optional Celeborn client") { + val conf = offHeapSQLConf() + conf.setConfString(CometConf.COMET_ENABLED.key, "true") conf.setConfString(CometConf.COMET_SHUFFLE_ENABLED.key, "true") conf.setConfString(CometConf.COMET_SHUFFLE_MODE.key, "native") conf.setConfString("spark.shuffle.manager", classOf[CometCelebornShuffleManager].getName) @@ -94,7 +118,7 @@ class CometSparkSessionExtensionsSuite extends CometTestBase { } test("the stock Celeborn manager cannot accept Comet shuffle dependencies") { - val conf = new SQLConf + val conf = offHeapSQLConf() conf.setConfString(CometConf.COMET_ENABLED.key, "true") conf.setConfString(CometConf.COMET_SHUFFLE_ENABLED.key, "true") conf.setConfString( diff --git a/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala b/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala index 4673740dafc..25ff179f24b 100644 --- a/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala @@ -23,9 +23,10 @@ import java.io.File import org.apache.logging.log4j.Level import org.apache.spark.sql.{CometTestBase, SaveMode} +import org.apache.spark.sql.comet.CometPlan import org.apache.spark.sql.internal.StaticSQLConf -import org.apache.comet.COMET_VERSION +import org.apache.comet.{COMET_VERSION, CometConf, CometSparkSessionExtensions} class CometPluginsSuite extends CometTestBase { override protected def sparkConf: SparkConf = { @@ -253,3 +254,45 @@ 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") { + // Logging derives the logger name by stripping the object's trailing '$' + val logger = CometSparkSessionExtensions.getClass.getName.stripSuffix("$") + val appender = new LogAppender("off-heap mode warning") + withParquetTable((0 until 10).map(i => (i, i.toString)), "tbl") { + withLogAppender(appender, Seq(logger), 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("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) + } + } + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/CometTPCQueryBase.scala b/spark/src/test/scala/org/apache/spark/sql/CometTPCQueryBase.scala index a3437559757..1cf39c3a248 100644 --- a/spark/src/test/scala/org/apache/spark/sql/CometTPCQueryBase.scala +++ b/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() diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBenchmarkBase.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBenchmarkBase.scala index 708e9de02e2..385c1eb1f93 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBenchmarkBase.scala +++ b/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() diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastHashJoinBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastHashJoinBenchmark.scala index 152d79dbb98..da1a857f588 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastHashJoinBenchmark.scala +++ b/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() diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastNestedLoopJoinBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastNestedLoopJoinBenchmark.scala index f8fe6100d4d..c1cfab60fcb 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastNestedLoopJoinBenchmark.scala +++ b/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() diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometExecBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometExecBenchmark.scala index 2e7e143a472..baa3306238c 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometExecBenchmark.scala +++ b/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 diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashJoinBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashJoinBenchmark.scala index ef2c325c547..f339ef7aa72 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashJoinBenchmark.scala +++ b/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() diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleBenchmark.scala index 64d9d0237af..17d96571c14 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleBenchmark.scala +++ b/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 diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometSortMergeJoinBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometSortMergeJoinBenchmark.scala index a429da58350..c0f976563a7 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometSortMergeJoinBenchmark.scala +++ b/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() From 2f207b0f2872f2f674f27ec139c3747aca01d752 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 24 Sep 2026 09:40:29 -0600 Subject: [PATCH 2/4] fix: read the off-heap setting from the SparkContext's conf When the SparkContext already exists, SparkSession.Builder copies core configs such as spark.memory.offHeap.enabled into the session's SQLConf without applying them, so a session could report off-heap memory while the executors ran on-heap. isCometLoaded now reads the setting from the SparkContext's conf, which is what the executors use. spark.comet.exec.onHeap.enabled is still read from the session conf. The unit test that set spark.memory.offHeap.enabled on a bare SQLConf no longer applies, so CometSparkSessionExtensionsSuite is back to its original form and CometPluginsExtensionOnlySuite covers the builder path. --- .../comet/CometSparkSessionExtensions.scala | 14 ++++---- .../CometSparkSessionExtensionsSuite.scala | 32 +++---------------- .../org/apache/spark/CometPluginsSuite.scala | 19 ++++++++++- 3 files changed, 28 insertions(+), 37 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala index b48ad248e6a..35750375ab4 100644 --- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala +++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala @@ -131,8 +131,12 @@ object CometSparkSessionExtensions extends Logging { } // CometDriverPlugin makes the same check before registering this extension, but an - // application can also register the extension directly with spark.sql.extensions. - if (!isOffHeapEnabled(conf) && !COMET_ONHEAP_ENABLED.get(conf)) { + // 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 } @@ -299,12 +303,6 @@ object CometSparkSessionExtensions extends Logging { sparkConf.getBoolean("spark.memory.offHeap.enabled", false) } - // Spark copies the SparkConf into each session's SQLConf and, by default, refuses to set core - // configs such as this one at session level, so this is the application's memory mode. - private def isOffHeapEnabled(conf: SQLConf): Boolean = { - conf.getConfString("spark.memory.offHeap.enabled", "false").toBoolean - } - /** * Record a fallback reason on a `TreeNode` (a Spark operator or expression) explaining why * Comet cannot accelerate it. Reasons recorded here are surfaced in extended explain output diff --git a/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala b/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala index 334a4383dfd..331d4b4ece9 100644 --- a/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala @@ -36,16 +36,8 @@ class CometSparkSessionExtensionsSuite extends CometTestBase { import CometSparkSessionExtensions._ - // isCometLoaded also requires off-heap memory, so the tests of its other checks start from a - // conf that has it. - private def offHeapSQLConf(): SQLConf = { - val conf = new SQLConf - conf.setConfString("spark.memory.offHeap.enabled", "true") - conf - } - test("isCometLoaded") { - val conf = offHeapSQLConf() + val conf = new SQLConf // Disable Comet shuffle so this test can focus on other checks without needing // spark.shuffle.manager to be set. conf.setConfString(CometConf.COMET_SHUFFLE_ENABLED.key, "false") @@ -71,7 +63,7 @@ class CometSparkSessionExtensionsSuite extends CometTestBase { } test("isCometLoaded requires CometShuffleManager when shuffle.enabled=true") { - val conf = offHeapSQLConf() + val conf = new SQLConf conf.setConfString(CometConf.COMET_ENABLED.key, "true") // Default: shuffle.enabled=true. Without spark.shuffle.manager set, Comet must be disabled. @@ -89,24 +81,8 @@ class CometSparkSessionExtensionsSuite extends CometTestBase { assert(isCometLoaded(conf)) } - test("isCometLoaded requires off-heap memory unless on-heap mode is enabled") { - val conf = new SQLConf - conf.setConfString(CometConf.COMET_ENABLED.key, "true") - conf.setConfString(CometConf.COMET_SHUFFLE_ENABLED.key, "false") - // Set explicitly, since ENABLE_COMET_ONHEAP in the environment changes the default. - conf.setConfString(CometConf.COMET_ONHEAP_ENABLED.key, "false") - assert(!isCometLoaded(conf)) - - conf.setConfString("spark.memory.offHeap.enabled", "true") - assert(isCometLoaded(conf)) - - conf.setConfString("spark.memory.offHeap.enabled", "false") - conf.setConfString(CometConf.COMET_ONHEAP_ENABLED.key, "true") - assert(isCometLoaded(conf)) - } - test("the composite manager is recognized without requiring the optional Celeborn client") { - val conf = offHeapSQLConf() + val conf = new SQLConf conf.setConfString(CometConf.COMET_ENABLED.key, "true") conf.setConfString(CometConf.COMET_SHUFFLE_ENABLED.key, "true") conf.setConfString(CometConf.COMET_SHUFFLE_MODE.key, "native") @@ -118,7 +94,7 @@ class CometSparkSessionExtensionsSuite extends CometTestBase { } test("the stock Celeborn manager cannot accept Comet shuffle dependencies") { - val conf = offHeapSQLConf() + val conf = new SQLConf conf.setConfString(CometConf.COMET_ENABLED.key, "true") conf.setConfString(CometConf.COMET_SHUFFLE_ENABLED.key, "true") conf.setConfString( diff --git a/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala b/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala index 25ff179f24b..189c2ef3002 100644 --- a/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala @@ -22,7 +22,7 @@ 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.internal.StaticSQLConf @@ -287,6 +287,23 @@ class CometPluginsExtensionOnlySuite extends CometTestBase { .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") { From 5c41a47257dacaf0118d5654cf45f059c1cb46a9 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 24 Sep 2026 10:00:17 -0600 Subject: [PATCH 3/4] fix: check the running shuffle manager instead of the session conf isCometShuffleManagerEnabled read spark.shuffle.manager from the session's SQLConf. When the SparkContext already exists, SparkSession.Builder copies core configs into the SQLConf without applying them, so a session could name CometShuffleManager while Spark's own manager ran the shuffle. Comet then planned Comet shuffles that failed at read time with a ClassCastException. It now inspects the manager that SparkEnv holds. --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + .../comet/CometSparkSessionExtensions.scala | 18 ++++---- .../shuffle/CometShuffleExchangeExec.scala | 2 +- .../CometSparkSessionExtensionsSuite.scala | 44 ++++++------------- .../org/apache/spark/CometPluginsSuite.scala | 42 ++++++++++++++++++ .../CometCelebornShufflePlanningSuite.scala | 2 +- 7 files changed, 69 insertions(+), 41 deletions(-) diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index fa4167d1c58..8fb9b373d31 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -544,6 +544,7 @@ jobs: 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 diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index b33da808c64..018fb004013 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -192,6 +192,7 @@ jobs: 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 diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala index 35750375ab4..f274248a2b2 100644 --- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala +++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala @@ -141,7 +141,7 @@ object CometSparkSessionExtensions extends Logging { return false } - if (COMET_SHUFFLE_ENABLED.get(conf) && !isCometShuffleManagerEnabled(conf)) { + 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 " + @@ -185,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] = @@ -198,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 diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index c4fbf094cc2..c2dd19ebcd1 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/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) diff --git a/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala b/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala index 331d4b4ece9..fa40dbb3ab6 100644 --- a/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala +++ b/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", diff --git a/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala b/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala index 189c2ef3002..aa6db104a38 100644 --- a/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala @@ -24,6 +24,7 @@ import java.io.File import org.apache.logging.log4j.Level 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, CometConf, CometSparkSessionExtensions} @@ -313,3 +314,44 @@ class CometPluginsExtensionOnlySuite extends CometTestBase { } } } + +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) + } + } + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala index 9617ec5644d..0e2b47e15ba 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala +++ b/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)) From e04526949be0c2a4522f005d535ee8894da53732 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 24 Sep 2026 11:46:11 -0600 Subject: [PATCH 4/4] test: listen on the org.apache.comet logger in CometPluginsExtensionOnlySuite withLogAppender on a logger with no LoggerConfig of its own makes log4j create one, copying the root config's additivity, which is false. The config outlives the test, so after the off-heap warning test listened on CometSparkSessionExtensions directly, that logger's events never reached an appender on org.apache.comet again. CometExecRuleSuite, which runs later in the same JVM in both PR workflows, counts a warning that CometSparkSessionExtensions logs through such an appender, and found none. --- .../scala/org/apache/spark/CometPluginsSuite.scala | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala b/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala index aa6db104a38..ebe89496295 100644 --- a/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala @@ -27,7 +27,7 @@ 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, CometConf, CometSparkSessionExtensions} +import org.apache.comet.{COMET_VERSION, CometConf} class CometPluginsSuite extends CometTestBase { override protected def sparkConf: SparkConf = { @@ -274,11 +274,13 @@ class CometPluginsExtensionOnlySuite extends CometTestBase { private val query = "SELECT _1 FROM tbl WHERE _1 > 5" test("Comet is disabled when off-heap memory is disabled") { - // Logging derives the logger name by stripping the object's trailing '$' - val logger = CometSparkSessionExtensions.getClass.getName.stripSuffix("$") + // 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(logger), Some(Level.WARN)) { + withLogAppender(appender, Seq("org.apache.comet"), Some(Level.WARN)) { val (_, plan) = checkSparkAnswer(query) assert(collect(plan) { case op: CometPlan => op }.isEmpty, plan) }