From 2db53f3e9df609b1f2fec20807531dc7c3c028b6 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Fri, 30 Oct 2026 08:46:02 -0600
Subject: [PATCH 1/3] feat: set CometShuffleManager from CometPlugin on Spark
 4.0 and later

From Spark 4.0 (SPARK-45762) SparkContext creates the shuffle manager right
after it initializes the plugins, from the same SparkConf, and executors take
their conf from the driver's when they register. CometDriverPlugin can
therefore choose the manager for an application that did not, instead of
leaving Comet to disable itself.

An application that sets spark.shuffle.manager keeps its value. Before 4.0 the
SparkEnv already holds the manager when plugins initialize, so the plugin does
nothing there.
---
 docs/source/about/faq.md                      |  6 +-
 .../user-guide/latest/tuning/shuffle.md       |  5 +-
 .../scala/org/apache/comet/CometConf.scala    |  4 +-
 .../comet/CometSparkSessionExtensions.scala   |  1 +
 .../main/scala/org/apache/spark/Plugins.scala | 22 ++++++
 .../org/apache/spark/CometPluginsSuite.scala  | 71 ++++++++++++++++++-
 6 files changed, 102 insertions(+), 7 deletions(-)

diff --git a/docs/source/about/faq.md b/docs/source/about/faq.md
index 29c39910ed1..842d8bc77ef 100644
--- a/docs/source/about/faq.md
+++ b/docs/source/about/faq.md
@@ -127,9 +127,9 @@ Common causes of a small speedup are:
 - **Fallbacks inside a stage.** Each switch between Comet and Spark converts data between columnar
   and row formats, which can cost more than Comet saves. See
   [Reducing Row/Columnar Conversion Overhead](../user-guide/latest/tuning/transitions.md).
-- **Comet's shuffle manager is not set.** If `spark.shuffle.manager` does not name Comet's shuffle
-  manager, Comet disables itself and logs a warning, unless `spark.comet.shuffle.enabled=false`, in
-  which case Comet runs but every shuffle runs in Spark. See
+- **Comet's shuffle manager is not set.** On Spark 3.4 and 3.5, or when the application sets
+  `spark.shuffle.manager` to another manager, Comet disables itself and logs a warning, unless
+  `spark.comet.shuffle.enabled=false`, in which case Comet runs but every shuffle runs in Spark. See
   [Shuffle Tuning](../user-guide/latest/tuning/shuffle.md).
 - **Too little memory.** Native operators spill to disk when the off-heap pool is too small. See
   [Memory Tuning](../user-guide/latest/tuning/memory.md).
diff --git a/docs/source/user-guide/latest/tuning/shuffle.md b/docs/source/user-guide/latest/tuning/shuffle.md
index 7764bfa06d9..ce19f3af403 100644
--- a/docs/source/user-guide/latest/tuning/shuffle.md
+++ b/docs/source/user-guide/latest/tuning/shuffle.md
@@ -29,7 +29,10 @@ spark.comet.shuffle.enabled=true
 ```
 
 `spark.shuffle.manager` is a Spark static configuration which cannot be changed at runtime.
-It must be set before the Spark context is created. You can enable or disable Comet shuffle
+It must be set before the Spark context is created. On Spark 4.0 and later, `CometPlugin` sets it to
+`CometShuffleManager` if the application leaves it unset, so the line above can be omitted there. An
+application that sets `spark.shuffle.manager` itself keeps that value. On Spark 3.4 and 3.5, Spark creates
+the shuffle manager before it initializes plugins, so the application has to set it. You can enable or disable Comet shuffle
 at runtime by setting `spark.comet.shuffle.enabled` to `true` or `false`.
 Once it is disabled, the configured shuffle manager handles ordinary Spark shuffle dependencies
 without Comet's shuffle implementation.
diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala
index 6003703a630..b6e6133eef0 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -594,7 +594,9 @@ object CometConf extends ShimCometConf {
           "`org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManager`. " +
           "Celeborn requires explicit native mode and compatible application settings. " +
           "`spark.shuffle.manager` must be set before starting the Spark application and " +
-          "cannot be changed during the application.")
+          "cannot be changed during the application. On Spark 4.0 and later, " +
+          "`org.apache.spark.CometPlugin` sets it to `CometShuffleManager` when the " +
+          "application leaves it unset and `spark.comet.enabled` is true.")
       .booleanConf
       .createWithDefault(true)
 
diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
index 81ed81a7028..322ab3fb407 100644
--- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
+++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
@@ -170,6 +170,7 @@ object CometSparkSessionExtensions extends Logging {
         "Comet extension is disabled because spark.shuffle.manager is not set to " +
           s"${classOf[CometShuffleManager].getName} or " +
           s"${classOf[CometCelebornShuffleManager].getName}. " +
+          "CometPlugin sets it on Spark 4.0 and later when it is unset. " +
           "Comet provides limited benefit without its shuffle manager. " +
           s"Set ${COMET_SHUFFLE_ENABLED.key}=false to keep Comet enabled with " +
           "Spark's default shuffle manager.")
diff --git a/spark/src/main/scala/org/apache/spark/Plugins.scala b/spark/src/main/scala/org/apache/spark/Plugins.scala
index 70f3ce6cedc..a8460ee3baf 100644
--- a/spark/src/main/scala/org/apache/spark/Plugins.scala
+++ b/spark/src/main/scala/org/apache/spark/Plugins.scala
@@ -112,6 +112,8 @@ class CometDriverPlugin private[spark] (clock: Clock) extends DriverPlugin with
 
     val extraConfs = new ju.HashMap[String, String]()
 
+    // Before the cache serializer, which is installed only with one of Comet's shuffle managers.
+    CometDriverPlugin.maybeSetShuffleManager(sc.conf)
     CometDriverPlugin.maybeSetCacheSerializer(sc.conf, extraConfs)
     CometDriverPlugin.warnIfKryoRegistrationsMissing(sc.conf)
 
@@ -174,6 +176,26 @@ object CometDriverPlugin extends Logging {
   /** Spark config key under which the loaded Comet version is exposed at runtime. */
   val COMET_VERSION_CONFIG = "spark.comet.version"
 
+  // Choose Comet's shuffle manager for an application that did not choose one, so that Comet
+  // does not disable itself. From Spark 4.0 (SPARK-45762) SparkContext creates the shuffle manager
+  // right after it initializes the plugins, from the same SparkConf, and executors take their conf
+  // from the driver's when they register, which happens after this. Before 4.0 the SparkEnv holds
+  // the manager already, so setting the conf here would change nothing but the conf, and an
+  // application there has to set spark.shuffle.manager itself.
+  // An application that set spark.shuffle.manager keeps its choice, even Spark's own: it may
+  // depend on that manager, as Celeborn does, and Comet then disables itself as before.
+  private[apache] def maybeSetShuffleManager(conf: SparkConf): Unit = {
+    val key = "spark.shuffle.manager"
+    if (CometSparkSessionExtensions.isSpark40Plus &&
+      getBooleanConf(conf, CometConf.COMET_ENABLED) &&
+      getBooleanConf(conf, CometConf.COMET_SHUFFLE_ENABLED) &&
+      !conf.contains(key)) {
+      val value = classOf[CometShuffleManager].getName
+      conf.set(key, value)
+      logInfo(s"Auto-set $key=$value")
+    }
+  }
+
   // Use Comet's cache serializer only when the native in-memory cache scan can run, which needs
   // Comet and its native execution as well as the cache config. spark.sql.cache.serializer is
   // static, so an application that starts with Comet or native execution off would otherwise
diff --git a/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala b/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala
index f85045a57ec..48b999d9934 100644
--- a/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala
@@ -29,11 +29,11 @@ import org.apache.spark.api.plugin.{DriverPlugin, ExecutorPlugin, PluginContext,
 import org.apache.spark.scheduler.{SparkListenerApplicationEnd, SparkListenerEvent, SparkListenerExecutorMetricsUpdate, SparkListenerExecutorRemoved}
 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.comet.execution.shuffle.{CometCelebornShuffleManager, CometShuffleExchangeExec, CometShuffleManager}
 import org.apache.spark.sql.internal.StaticSQLConf
 import org.apache.spark.util.{JsonProtocol, ManualClock, Utils}
 
-import org.apache.comet.{COMET_VERSION, CometConf, CometExecIterator, CometExecutorMemoryUsage}
+import org.apache.comet.{COMET_VERSION, CometConf, CometExecIterator, CometExecutorMemoryUsage, CometSparkSessionExtensions}
 
 class CometPluginsSuite extends CometTestBase {
   override protected def sparkConf: SparkConf = {
@@ -328,6 +328,73 @@ class CometPluginsUnifiedModeSuite extends CometTestBase {
   }
 }
 
+// The application names no shuffle manager. From Spark 4.0 the plugin chooses Comet's, because
+// SparkContext creates the manager after it initializes the plugins. Before 4.0 the manager exists
+// by then, so the plugin cannot, and Comet disables itself.
+class CometPluginsDefaultShuffleManagerSuite extends CometTestBase {
+  override protected def sparkConf: SparkConf = {
+    val conf = new SparkConf()
+    conf.set("spark.plugins", "org.apache.spark.CometPlugin")
+    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("the plugin sets Comet's shuffle manager from Spark 4.0 and Comet runs") {
+    val cometManager = classOf[CometShuffleManager].getName
+    val expected = CometSparkSessionExtensions.isSpark40Plus
+    assert(spark.sparkContext.conf.contains("spark.shuffle.manager") == expected)
+    assert(SparkEnv.get.shuffleManager.isInstanceOf[CometShuffleManager] == expected)
+    if (expected) {
+      assert(spark.sparkContext.conf.get("spark.shuffle.manager") == cometManager)
+    }
+    withParquetTable((0 until 100).map(i => (i, (i % 7).toString)), "tbl") {
+      val (_, plan) = checkSparkAnswer(query)
+      assert(collect(plan) { case op: CometPlan => op }.nonEmpty == expected, plan)
+      assert(collect(plan) { case op: CometShuffleExchangeExec => op }.nonEmpty == expected, plan)
+    }
+  }
+}
+
+class CometDriverPluginShuffleManagerSuite extends SparkFunSuite {
+  private val key = "spark.shuffle.manager"
+
+  private def chosen(conf: SparkConf): Option[String] = {
+    CometDriverPlugin.maybeSetShuffleManager(conf)
+    conf.getOption(key)
+  }
+
+  private val cometManager = classOf[CometShuffleManager].getName
+
+  test("sets Comet's shuffle manager only on Spark 4.0 and later") {
+    val expected = if (CometSparkSessionExtensions.isSpark40Plus) Some(cometManager) else None
+    assert(chosen(new SparkConf(false)) == expected)
+  }
+
+  test("keeps the shuffle manager the application set") {
+    Seq("sort", "tungsten-sort", classOf[CometCelebornShuffleManager].getName, "com.example.Mgr")
+      .foreach { manager =>
+        assert(chosen(new SparkConf(false).set(key, manager)).contains(manager))
+      }
+  }
+
+  test("leaves the shuffle manager unset when Comet or Comet shuffle is disabled") {
+    Seq(CometConf.COMET_ENABLED, CometConf.COMET_SHUFFLE_ENABLED).foreach { entry =>
+      assert(chosen(new SparkConf(false).set(entry.key, "false")).isEmpty)
+    }
+  }
+
+  test("leaves the shuffle manager unset when a deprecated alternative disables Comet shuffle") {
+    val conf = new SparkConf(false)
+    conf.set("spark.comet.exec.shuffle.enabled", "false")
+    assert(chosen(conf).isEmpty)
+  }
+}
+
 class CometPluginsExtensionOnlySuite extends CometTestBase {
   // No plugin and no off-heap memory. CometTestBase registers CometSparkSessionExtensions
   // directly, as an application can with spark.sql.extensions.

From b630bbe444818f16581bce723855c42b04210003 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Fri, 30 Oct 2026 08:55:45 -0600
Subject: [PATCH 2/3] refactor: tighten the plugin's shuffle manager change and
 its docs

Use Spark's SHUFFLE_MANAGER key throughout, state the Spark 4.0 rule in fewer
places, add an upgrade guide entry, and run the plugin's real-session test in
CometPluginsUnifiedModeSuite rather than a session of its own.
---
 docs/source/about/faq.md                      |  6 +--
 .../user-guide/latest/migration-guide.md      |  9 ++++
 .../user-guide/latest/tuning/shuffle.md       |  8 ++--
 .../scala/org/apache/comet/CometConf.scala    |  3 +-
 .../comet/CometSparkSessionExtensions.scala   |  1 -
 .../main/scala/org/apache/spark/Plugins.scala | 23 ++++-----
 .../org/apache/spark/CometPluginsSuite.scala  | 48 +++++--------------
 7 files changed, 36 insertions(+), 62 deletions(-)

diff --git a/docs/source/about/faq.md b/docs/source/about/faq.md
index 842d8bc77ef..29c39910ed1 100644
--- a/docs/source/about/faq.md
+++ b/docs/source/about/faq.md
@@ -127,9 +127,9 @@ Common causes of a small speedup are:
 - **Fallbacks inside a stage.** Each switch between Comet and Spark converts data between columnar
   and row formats, which can cost more than Comet saves. See
   [Reducing Row/Columnar Conversion Overhead](../user-guide/latest/tuning/transitions.md).
-- **Comet's shuffle manager is not set.** On Spark 3.4 and 3.5, or when the application sets
-  `spark.shuffle.manager` to another manager, Comet disables itself and logs a warning, unless
-  `spark.comet.shuffle.enabled=false`, in which case Comet runs but every shuffle runs in Spark. See
+- **Comet's shuffle manager is not set.** If `spark.shuffle.manager` does not name Comet's shuffle
+  manager, Comet disables itself and logs a warning, unless `spark.comet.shuffle.enabled=false`, in
+  which case Comet runs but every shuffle runs in Spark. See
   [Shuffle Tuning](../user-guide/latest/tuning/shuffle.md).
 - **Too little memory.** Native operators spill to disk when the off-heap pool is too small. See
   [Memory Tuning](../user-guide/latest/tuning/memory.md).
diff --git a/docs/source/user-guide/latest/migration-guide.md b/docs/source/user-guide/latest/migration-guide.md
index feaa15aab06..d516b8f8393 100644
--- a/docs/source/user-guide/latest/migration-guide.md
+++ b/docs/source/user-guide/latest/migration-guide.md
@@ -57,6 +57,15 @@ Treat setting one of these keys as a temporary measure. If you find you cannot s
 legacy behavior, please open an issue describing your use case so it can be considered before the
 key is removed.
 
+## Upgrading to Comet 1.3.0
+
+### Shuffle Manager Set by the Plugin on Spark 4.0 and Later
+
+On Spark 4.0 and later, `CometPlugin` sets `spark.shuffle.manager` to `CometShuffleManager` when
+the application leaves it unset. Previously, an application that registered the plugin without a
+shuffle manager ran without Comet, with a warning. It now runs with Comet. To keep the previous
+behavior, set `spark.comet.enabled=false`, or set `spark.shuffle.manager` to the manager you want.
+
 ## Upgrading to Comet 1.2.0
 
 Comet `1.2.0` makes no behavior changes that need a `spark.comet.legacy.*` key. The changes below
diff --git a/docs/source/user-guide/latest/tuning/shuffle.md b/docs/source/user-guide/latest/tuning/shuffle.md
index ce19f3af403..49a8cc75b87 100644
--- a/docs/source/user-guide/latest/tuning/shuffle.md
+++ b/docs/source/user-guide/latest/tuning/shuffle.md
@@ -29,11 +29,9 @@ spark.comet.shuffle.enabled=true
 ```
 
 `spark.shuffle.manager` is a Spark static configuration which cannot be changed at runtime.
-It must be set before the Spark context is created. On Spark 4.0 and later, `CometPlugin` sets it to
-`CometShuffleManager` if the application leaves it unset, so the line above can be omitted there. An
-application that sets `spark.shuffle.manager` itself keeps that value. On Spark 3.4 and 3.5, Spark creates
-the shuffle manager before it initializes plugins, so the application has to set it. You can enable or disable Comet shuffle
-at runtime by setting `spark.comet.shuffle.enabled` to `true` or `false`.
+It must be set before the Spark context is created, except that on Spark 4.0 and later `CometPlugin`
+sets it to `CometShuffleManager` if the application leaves it unset. You can enable or disable
+Comet shuffle at runtime by setting `spark.comet.shuffle.enabled` to `true` or `false`.
 Once it is disabled, the configured shuffle manager handles ordinary Spark shuffle dependencies
 without Comet's shuffle implementation.
 
diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala
index b6e6133eef0..31c37611e08 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -595,8 +595,7 @@ object CometConf extends ShimCometConf {
           "Celeborn requires explicit native mode and compatible application settings. " +
           "`spark.shuffle.manager` must be set before starting the Spark application and " +
           "cannot be changed during the application. On Spark 4.0 and later, " +
-          "`org.apache.spark.CometPlugin` sets it to `CometShuffleManager` when the " +
-          "application leaves it unset and `spark.comet.enabled` is true.")
+          "`org.apache.spark.CometPlugin` sets it to `CometShuffleManager` if it is unset.")
       .booleanConf
       .createWithDefault(true)
 
diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
index 322ab3fb407..81ed81a7028 100644
--- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
+++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
@@ -170,7 +170,6 @@ object CometSparkSessionExtensions extends Logging {
         "Comet extension is disabled because spark.shuffle.manager is not set to " +
           s"${classOf[CometShuffleManager].getName} or " +
           s"${classOf[CometCelebornShuffleManager].getName}. " +
-          "CometPlugin sets it on Spark 4.0 and later when it is unset. " +
           "Comet provides limited benefit without its shuffle manager. " +
           s"Set ${COMET_SHUFFLE_ENABLED.key}=false to keep Comet enabled with " +
           "Spark's default shuffle manager.")
diff --git a/spark/src/main/scala/org/apache/spark/Plugins.scala b/spark/src/main/scala/org/apache/spark/Plugins.scala
index a8460ee3baf..f0b02f8c720 100644
--- a/spark/src/main/scala/org/apache/spark/Plugins.scala
+++ b/spark/src/main/scala/org/apache/spark/Plugins.scala
@@ -28,7 +28,7 @@ import scala.util.Try
 
 import org.apache.spark.api.plugin.{DriverPlugin, ExecutorPlugin, PluginContext, SparkPlugin}
 import org.apache.spark.internal.Logging
-import org.apache.spark.internal.config.{EVENT_LOG_ENABLED, EXECUTOR_MEMORY_OVERHEAD, EXECUTOR_MEMORY_OVERHEAD_FACTOR}
+import org.apache.spark.internal.config.{EVENT_LOG_ENABLED, EXECUTOR_MEMORY_OVERHEAD, EXECUTOR_MEMORY_OVERHEAD_FACTOR, SHUFFLE_MANAGER}
 import org.apache.spark.scheduler.{SparkListener, SparkListenerApplicationEnd, SparkListenerExecutorMetricsUpdate, SparkListenerExecutorRemoved}
 import org.apache.spark.serializer.KryoSerializer
 import org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer
@@ -176,23 +176,18 @@ object CometDriverPlugin extends Logging {
   /** Spark config key under which the loaded Comet version is exposed at runtime. */
   val COMET_VERSION_CONFIG = "spark.comet.version"
 
-  // Choose Comet's shuffle manager for an application that did not choose one, so that Comet
-  // does not disable itself. From Spark 4.0 (SPARK-45762) SparkContext creates the shuffle manager
-  // right after it initializes the plugins, from the same SparkConf, and executors take their conf
-  // from the driver's when they register, which happens after this. Before 4.0 the SparkEnv holds
-  // the manager already, so setting the conf here would change nothing but the conf, and an
-  // application there has to set spark.shuffle.manager itself.
-  // An application that set spark.shuffle.manager keeps its choice, even Spark's own: it may
-  // depend on that manager, as Celeborn does, and Comet then disables itself as before.
+  // Choose Comet's shuffle manager for an application that did not, so that Comet does not
+  // disable itself. Only from Spark 4.0 (SPARK-45762) does SparkContext create the manager after
+  // it initializes the plugins, from the same conf that executors then fetch. An application's own
+  // choice, even Spark's, is kept: it may depend on that manager, as Celeborn does.
   private[apache] def maybeSetShuffleManager(conf: SparkConf): Unit = {
-    val key = "spark.shuffle.manager"
     if (CometSparkSessionExtensions.isSpark40Plus &&
       getBooleanConf(conf, CometConf.COMET_ENABLED) &&
       getBooleanConf(conf, CometConf.COMET_SHUFFLE_ENABLED) &&
-      !conf.contains(key)) {
+      !conf.contains(SHUFFLE_MANAGER)) {
       val value = classOf[CometShuffleManager].getName
-      conf.set(key, value)
-      logInfo(s"Auto-set $key=$value")
+      conf.set(SHUFFLE_MANAGER.key, value)
+      logInfo(s"Auto-set ${SHUFFLE_MANAGER.key}=$value")
     }
   }
 
@@ -286,7 +281,7 @@ object CometDriverPlugin extends Logging {
   // class name.
   private def isCometShuffleManager(conf: SparkConf): Boolean =
     Set(classOf[CometShuffleManager].getName, classOf[CometCelebornShuffleManager].getName)
-      .contains(conf.get("spark.shuffle.manager", ""))
+      .contains(conf.get(SHUFFLE_MANAGER.key, ""))
 
   // Comet's native allocations are made by the Rust global allocator and live in the native heap.
   // In off-heap mode the share that operators reserve is charged against a memory pool, but
diff --git a/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala b/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala
index 48b999d9934..2f3bebafa80 100644
--- a/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala
@@ -26,6 +26,7 @@ import java.util.concurrent.TimeUnit
 
 import org.apache.logging.log4j.Level
 import org.apache.spark.api.plugin.{DriverPlugin, ExecutorPlugin, PluginContext, SparkPlugin}
+import org.apache.spark.internal.config.SHUFFLE_MANAGER
 import org.apache.spark.scheduler.{SparkListenerApplicationEnd, SparkListenerEvent, SparkListenerExecutorMetricsUpdate, SparkListenerExecutorRemoved}
 import org.apache.spark.sql.{CometTestBase, SaveMode, SparkSession}
 import org.apache.spark.sql.comet.CometPlan
@@ -326,50 +327,27 @@ class CometPluginsUnifiedModeSuite extends CometTestBase {
     assert(execMemOverhead3 == "1G")
     assert(execMemOverhead4 == "1G")
   }
-}
-
-// The application names no shuffle manager. From Spark 4.0 the plugin chooses Comet's, because
-// SparkContext creates the manager after it initializes the plugins. Before 4.0 the manager exists
-// by then, so the plugin cannot, and Comet disables itself.
-class CometPluginsDefaultShuffleManagerSuite extends CometTestBase {
-  override protected def sparkConf: SparkConf = {
-    val conf = new SparkConf()
-    conf.set("spark.plugins", "org.apache.spark.CometPlugin")
-    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"
 
+  // This suite names no shuffle manager. From Spark 4.0 the plugin chooses Comet's; before that
+  // Spark has built its own by the time the plugin runs, and Comet disables itself.
   test("the plugin sets Comet's shuffle manager from Spark 4.0 and Comet runs") {
-    val cometManager = classOf[CometShuffleManager].getName
     val expected = CometSparkSessionExtensions.isSpark40Plus
-    assert(spark.sparkContext.conf.contains("spark.shuffle.manager") == expected)
     assert(SparkEnv.get.shuffleManager.isInstanceOf[CometShuffleManager] == expected)
-    if (expected) {
-      assert(spark.sparkContext.conf.get("spark.shuffle.manager") == cometManager)
-    }
     withParquetTable((0 until 100).map(i => (i, (i % 7).toString)), "tbl") {
-      val (_, plan) = checkSparkAnswer(query)
-      assert(collect(plan) { case op: CometPlan => op }.nonEmpty == expected, plan)
+      val (_, plan) = checkSparkAnswer("SELECT _2, count(*) FROM tbl GROUP BY _2")
       assert(collect(plan) { case op: CometShuffleExchangeExec => op }.nonEmpty == expected, plan)
     }
   }
 }
 
 class CometDriverPluginShuffleManagerSuite extends SparkFunSuite {
-  private val key = "spark.shuffle.manager"
+  private val cometManager = classOf[CometShuffleManager].getName
 
   private def chosen(conf: SparkConf): Option[String] = {
     CometDriverPlugin.maybeSetShuffleManager(conf)
-    conf.getOption(key)
+    conf.getOption(SHUFFLE_MANAGER.key)
   }
 
-  private val cometManager = classOf[CometShuffleManager].getName
-
   test("sets Comet's shuffle manager only on Spark 4.0 and later") {
     val expected = if (CometSparkSessionExtensions.isSpark40Plus) Some(cometManager) else None
     assert(chosen(new SparkConf(false)) == expected)
@@ -378,21 +356,17 @@ class CometDriverPluginShuffleManagerSuite extends SparkFunSuite {
   test("keeps the shuffle manager the application set") {
     Seq("sort", "tungsten-sort", classOf[CometCelebornShuffleManager].getName, "com.example.Mgr")
       .foreach { manager =>
-        assert(chosen(new SparkConf(false).set(key, manager)).contains(manager))
+        assert(chosen(new SparkConf(false).set(SHUFFLE_MANAGER.key, manager)).contains(manager))
       }
   }
 
   test("leaves the shuffle manager unset when Comet or Comet shuffle is disabled") {
-    Seq(CometConf.COMET_ENABLED, CometConf.COMET_SHUFFLE_ENABLED).foreach { entry =>
-      assert(chosen(new SparkConf(false).set(entry.key, "false")).isEmpty)
+    val disablingKeys = Seq(CometConf.COMET_ENABLED.key, CometConf.COMET_SHUFFLE_ENABLED.key) ++
+      CometConf.COMET_SHUFFLE_ENABLED.alternatives
+    disablingKeys.foreach { key =>
+      assert(chosen(new SparkConf(false).set(key, "false")).isEmpty, key)
     }
   }
-
-  test("leaves the shuffle manager unset when a deprecated alternative disables Comet shuffle") {
-    val conf = new SparkConf(false)
-    conf.set("spark.comet.exec.shuffle.enabled", "false")
-    assert(chosen(conf).isEmpty)
-  }
 }
 
 class CometPluginsExtensionOnlySuite extends CometTestBase {

From bfc14c7671cda6303ef0f0464f581d11da37f540 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Wed, 4 Nov 2026 14:32:37 -0700
Subject: [PATCH 3/3] fix: leave the shuffle manager unset where Comet disables
 itself, and update docs

The plugin no longer sets CometShuffleManager on a big-endian host, with
spark.sql.parquet.int96TimestampConversion=true, or when the native library
fails to load, since Comet would not run and Spark's exchanges would copy
every row. isCometLoaded and the plugin share one helper for these checks.

Register CometDriverPluginShuffleManagerSuite in the PR build workflows, and
bring plugin_overview.md, the FAQ and the upgrade guide up to date.
---
 .github/workflows/pr_build_linux.yml          |  1 +
 .github/workflows/pr_build_macos.yml          |  1 +
 docs/source/about/faq.md                      |  4 ++-
 .../contributor-guide/plugin_overview.md      | 11 +++++-
 .../user-guide/latest/migration-guide.md      |  9 +++++
 .../comet/CometSparkSessionExtensions.scala   | 36 ++++++++++++++-----
 .../main/scala/org/apache/spark/Plugins.scala | 19 ++++++++--
 .../org/apache/spark/CometPluginsSuite.scala  |  7 +++-
 8 files changed, 74 insertions(+), 14 deletions(-)

diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml
index 52e86f59ffd..99bac1712c6 100644
--- a/.github/workflows/pr_build_linux.yml
+++ b/.github/workflows/pr_build_linux.yml
@@ -589,6 +589,7 @@ jobs:
               org.apache.spark.CometPluginsUnifiedModeSuite
               org.apache.spark.CometPluginsExtensionOnlySuite
               org.apache.spark.CometPluginsSparkShuffleManagerSuite
+              org.apache.spark.CometDriverPluginShuffleManagerSuite
               org.apache.spark.CometPluginsEventLogSuite
               org.apache.comet.rules.CometScanRuleSuite
               org.apache.comet.rules.CometScanContribSuite
diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml
index b1d89ec1240..ad7bcc134b3 100644
--- a/.github/workflows/pr_build_macos.yml
+++ b/.github/workflows/pr_build_macos.yml
@@ -255,6 +255,7 @@ jobs:
               org.apache.spark.CometPluginsUnifiedModeSuite
               org.apache.spark.CometPluginsExtensionOnlySuite
               org.apache.spark.CometPluginsSparkShuffleManagerSuite
+              org.apache.spark.CometDriverPluginShuffleManagerSuite
               org.apache.spark.CometPluginsEventLogSuite
               org.apache.comet.rules.CometScanRuleSuite
               org.apache.comet.rules.CometScanContribSuite
diff --git a/docs/source/about/faq.md b/docs/source/about/faq.md
index 29c39910ed1..3276fb1ac40 100644
--- a/docs/source/about/faq.md
+++ b/docs/source/about/faq.md
@@ -129,7 +129,9 @@ Common causes of a small speedup are:
   [Reducing Row/Columnar Conversion Overhead](../user-guide/latest/tuning/transitions.md).
 - **Comet's shuffle manager is not set.** If `spark.shuffle.manager` does not name Comet's shuffle
   manager, Comet disables itself and logs a warning, unless `spark.comet.shuffle.enabled=false`, in
-  which case Comet runs but every shuffle runs in Spark. See
+  which case Comet runs but every shuffle runs in Spark. On Spark 4.0 and later, `CometPlugin` sets
+  the manager when the application leaves `spark.shuffle.manager` unset, so this applies there only
+  when the application names another manager. On Spark 3.4 and 3.5 the application must set it. See
   [Shuffle Tuning](../user-guide/latest/tuning/shuffle.md).
 - **Too little memory.** Native operators spill to disk when the off-heap pool is too small. See
   [Memory Tuning](../user-guide/latest/tuning/memory.md).
diff --git a/docs/source/contributor-guide/plugin_overview.md b/docs/source/contributor-guide/plugin_overview.md
index 072f59e01a2..46787e4a75f 100644
--- a/docs/source/contributor-guide/plugin_overview.md
+++ b/docs/source/contributor-guide/plugin_overview.md
@@ -46,6 +46,13 @@ version. If neither `spark.memory.offHeap.enabled` nor `spark.comet.exec.onHeap.
 and skips the remaining steps. Otherwise it:
 
 - Appends `CometSparkSessionExtensions` to `spark.sql.extensions`, unless it is already listed.
+- On Spark 4.0 and later, sets `spark.shuffle.manager` to Comet's `CometShuffleManager`, when
+  `spark.comet.enabled` and Comet shuffle are on and the application has not set `spark.shuffle.manager`. An
+  application's own value is always kept. It is not set where Comet would disable itself anyway, such as on a
+  big-endian platform, with `spark.sql.parquet.int96TimestampConversion=true`, or when the native library fails to
+  load. Before Spark 4.0 the plugin cannot do this, because Spark has created the shuffle manager by the time plugins
+  run. It runs before the cache serializer step, since that serializer is installed only with one of Comet's shuffle
+  managers.
 - Sets `spark.sql.cache.serializer` to Comet's `ArrowCachedBatchSerializer` when
   `spark.comet.exec.inMemoryCache.enabled=true`, unless the application has chosen a different serializer.
   It leaves Spark's serializer in place where Comet could not scan its format natively or Kryo would
@@ -139,7 +146,9 @@ executing, the resulting Arrow batches are imported into the JVM using Arrow FFI
 ## Shuffle
 
 Comet integrates with Spark's shuffle mechanism, optimizing both shuffle writes and reads. Comet's shuffle manager
-must be registered with Spark using the following configuration:
+must be registered with Spark using the following configuration. On Spark 4.0 and later, `CometDriverPlugin`
+registers it when the application leaves `spark.shuffle.manager` unset (see above); on earlier versions, the
+application must set it:
 
 ```
 --conf spark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager
diff --git a/docs/source/user-guide/latest/migration-guide.md b/docs/source/user-guide/latest/migration-guide.md
index d516b8f8393..e03cc17ea56 100644
--- a/docs/source/user-guide/latest/migration-guide.md
+++ b/docs/source/user-guide/latest/migration-guide.md
@@ -59,6 +59,9 @@ key is removed.
 
 ## Upgrading to Comet 1.3.0
 
+Comet `1.3.0` makes no behavior changes that need a `spark.comet.legacy.*` key. The change below
+needs none either, but check whether it applies to your deployment.
+
 ### Shuffle Manager Set by the Plugin on Spark 4.0 and Later
 
 On Spark 4.0 and later, `CometPlugin` sets `spark.shuffle.manager` to `CometShuffleManager` when
@@ -66,6 +69,12 @@ the application leaves it unset. Previously, an application that registered the
 shuffle manager ran without Comet, with a warning. It now runs with Comet. To keep the previous
 behavior, set `spark.comet.enabled=false`, or set `spark.shuffle.manager` to the manager you want.
 
+These applications never ran Comet, so they may have left `spark.executor.memoryOverhead` at Spark's
+default. Comet's native memory outside the off-heap pool comes out of that overhead, and YARN or
+Kubernetes can kill an executor that runs out of it. The plugin warns at startup when it is unset,
+but check it before you upgrade; see
+[Configuring Comet Memory](tuning/memory.md#configuring-comet-memory).
+
 ## Upgrading to Comet 1.2.0
 
 Comet `1.2.0` makes no behavior changes that need a `spark.comet.legacy.*` key. The changes below
diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
index 81ed81a7028..de923f45511 100644
--- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
+++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
@@ -145,10 +145,6 @@ object CometSparkSessionExtensions extends Logging {
    * Checks whether Comet extension should be loaded for Spark.
    */
   private[comet] def isCometLoaded(conf: SQLConf): Boolean = {
-    if (isBigEndian) {
-      logInfo("Comet extension is disabled because platform is big-endian")
-      return false
-    }
     if (!COMET_ENABLED.get(conf)) {
       logInfo(s"Comet extension is disabled, please turn on ${COMET_ENABLED.key} to enable it")
       return false
@@ -176,21 +172,43 @@ object CometSparkSessionExtensions extends Logging {
       return false
     }
 
+    unsupportedEnvironmentReason(
+      conf.getConf(SQLConf.PARQUET_INT96_TIMESTAMP_CONVERSION),
+      COMET_NATIVE_LOAD_REQUIRED.get(conf)).isEmpty
+  }
+
+  /**
+   * Why Comet disables itself on this platform or with these settings, whatever shuffle manager
+   * the application uses, or None if nothing rules it out. [[isCometLoaded]] applies it, and so
+   * does CometDriverPlugin before it chooses a shuffle manager, so that an application where
+   * Comet cannot run keeps Spark's.
+   *
+   * This loads the native library on demand. If that fails and `nativeLoadRequired` is set, it
+   * throws instead of returning a reason.
+   */
+  private[apache] def unsupportedEnvironmentReason(
+      int96TimestampConversion: Boolean,
+      nativeLoadRequired: Boolean): Option[String] = {
+    if (isBigEndian) {
+      logInfo("Comet extension is disabled because platform is big-endian")
+      return Some("the platform is big-endian")
+    }
+
     // We don't support INT96 timestamps written by Apache Impala in a different timezone yet
-    if (conf.getConf(SQLConf.PARQUET_INT96_TIMESTAMP_CONVERSION)) {
+    if (int96TimestampConversion) {
       logWarning(
         "Comet extension is disabled, because it currently doesn't support" +
           s" ${SQLConf.PARQUET_INT96_TIMESTAMP_CONVERSION} setting to true.")
-      return false
+      return Some(s"${SQLConf.PARQUET_INT96_TIMESTAMP_CONVERSION.key} is true")
     }
 
     try {
       // This will load the Comet native lib on demand, and if success, should set
       // `NativeBase.loaded` to true
-      NativeBase.isLoaded
+      if (NativeBase.isLoaded) None else Some("the native library is not loaded")
     } catch {
       case e: Throwable =>
-        if (COMET_NATIVE_LOAD_REQUIRED.get(conf)) {
+        if (nativeLoadRequired) {
           throw new CometRuntimeException(
             "Error when loading native library. Please fix the error and try again, or fallback " +
               s"to Spark by setting ${COMET_ENABLED.key} to false",
@@ -201,7 +219,7 @@ object CometSparkSessionExtensions extends Logging {
               "Falling back to Spark",
             e)
         }
-        false
+        Some("the native library failed to load")
     }
   }
 
diff --git a/spark/src/main/scala/org/apache/spark/Plugins.scala b/spark/src/main/scala/org/apache/spark/Plugins.scala
index f0b02f8c720..f6d54fbc81f 100644
--- a/spark/src/main/scala/org/apache/spark/Plugins.scala
+++ b/spark/src/main/scala/org/apache/spark/Plugins.scala
@@ -33,7 +33,7 @@ import org.apache.spark.scheduler.{SparkListener, SparkListenerApplicationEnd, S
 import org.apache.spark.serializer.KryoSerializer
 import org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer
 import org.apache.spark.sql.comet.execution.shuffle.{CometCelebornShuffleManager, CometShuffleManager}
-import org.apache.spark.sql.internal.StaticSQLConf
+import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf}
 import org.apache.spark.util.{Clock, SystemClock}
 
 import org.apache.comet.{COMET_VERSION, CometExecIterator, CometExecutorMemoryUsage, CometSparkSessionExtensions, NativeBase}
@@ -180,17 +180,32 @@ object CometDriverPlugin extends Logging {
   // disable itself. Only from Spark 4.0 (SPARK-45762) does SparkContext create the manager after
   // it initializes the plugins, from the same conf that executors then fetch. An application's own
   // choice, even Spark's, is kept: it may depend on that manager, as Celeborn does.
+  // Nor is it set where Comet disables itself whatever the manager, since Comet would not run and
+  // Spark would copy every row of its exchanges for want of a SortShuffleManager.
   private[apache] def maybeSetShuffleManager(conf: SparkConf): Unit = {
     if (CometSparkSessionExtensions.isSpark40Plus &&
       getBooleanConf(conf, CometConf.COMET_ENABLED) &&
       getBooleanConf(conf, CometConf.COMET_SHUFFLE_ENABLED) &&
-      !conf.contains(SHUFFLE_MANAGER)) {
+      !conf.contains(SHUFFLE_MANAGER) &&
+      !cometUnsupportedAtStartup(conf)) {
       val value = classOf[CometShuffleManager].getName
       conf.set(SHUFFLE_MANAGER.key, value)
       logInfo(s"Auto-set ${SHUFFLE_MANAGER.key}=$value")
     }
   }
 
+  // Applies the checks that isCometLoaded makes on the platform and settings. A library that
+  // fails to load when nativeLoadRequired is set counts as supported: isCometLoaded throws then,
+  // when Comet first plans a query, and the application should not fail here instead.
+  private def cometUnsupportedAtStartup(conf: SparkConf): Boolean = {
+    val int96Conversion = conf.getBoolean(SQLConf.PARQUET_INT96_TIMESTAMP_CONVERSION.key, false)
+    val nativeLoadRequired = getBooleanConf(conf, CometConf.COMET_NATIVE_LOAD_REQUIRED)
+    Try(
+      CometSparkSessionExtensions
+        .unsupportedEnvironmentReason(int96Conversion, nativeLoadRequired)
+        .isDefined).getOrElse(false)
+  }
+
   // Use Comet's cache serializer only when the native in-memory cache scan can run, which needs
   // Comet and its native execution as well as the cache config. spark.sql.cache.serializer is
   // static, so an application that starts with Comet or native execution off would otherwise
diff --git a/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala b/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala
index 2f3bebafa80..c4825a19d46 100644
--- a/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala
@@ -31,7 +31,7 @@ import org.apache.spark.scheduler.{SparkListenerApplicationEnd, SparkListenerEve
 import org.apache.spark.sql.{CometTestBase, SaveMode, SparkSession}
 import org.apache.spark.sql.comet.CometPlan
 import org.apache.spark.sql.comet.execution.shuffle.{CometCelebornShuffleManager, CometShuffleExchangeExec, CometShuffleManager}
-import org.apache.spark.sql.internal.StaticSQLConf
+import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf}
 import org.apache.spark.util.{JsonProtocol, ManualClock, Utils}
 
 import org.apache.comet.{COMET_VERSION, CometConf, CometExecIterator, CometExecutorMemoryUsage, CometSparkSessionExtensions}
@@ -367,6 +367,11 @@ class CometDriverPluginShuffleManagerSuite extends SparkFunSuite {
       assert(chosen(new SparkConf(false).set(key, "false")).isEmpty, key)
     }
   }
+
+  test("leaves the shuffle manager unset where Comet disables itself whatever the manager") {
+    val conf = new SparkConf(false).set(SQLConf.PARQUET_INT96_TIMESTAMP_CONVERSION.key, "true")
+    assert(chosen(conf).isEmpty)
+  }
 }
 
 class CometPluginsExtensionOnlySuite extends CometTestBase {