feat: set CometShuffleManager from CometPlugin on Spark 4.0 and later

.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

.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

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).

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

docs/source/user-guide/latest/migration-guide.md

@@ -57,6 +57,24 @@ 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
+
+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
+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.
+
+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

docs/source/user-guide/latest/tuning/shuffle.md

@@ -29,8 +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. 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.
 

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

@@ -594,7 +594,8 @@ 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` if it is unset.")
       .booleanConf
       .createWithDefault(true)
 

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")
     }
   }
 

spark/src/main/scala/org/apache/spark/Plugins.scala

@@ -28,12 +28,12 @@ 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
 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}
@@ -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,36 @@ 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, 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.
+  // 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) &&
+      !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
@@ -264,7 +296,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

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

@@ -26,14 +26,15 @@ 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
-import org.apache.spark.sql.comet.execution.shuffle.{CometShuffleExchangeExec, CometShuffleManager}
-import org.apache.spark.sql.internal.StaticSQLConf
+import org.apache.spark.sql.comet.execution.shuffle.{CometCelebornShuffleManager, CometShuffleExchangeExec, CometShuffleManager}
+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}
+import org.apache.comet.{COMET_VERSION, CometConf, CometExecIterator, CometExecutorMemoryUsage, CometSparkSessionExtensions}
 
 class CometPluginsSuite extends CometTestBase {
   override protected def sparkConf: SparkConf = {
@@ -326,6 +327,51 @@ class CometPluginsUnifiedModeSuite extends CometTestBase {
     assert(execMemOverhead3 == "1G")
     assert(execMemOverhead4 == "1G")
   }
+
+  // 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 expected = CometSparkSessionExtensions.isSpark40Plus
+    assert(SparkEnv.get.shuffleManager.isInstanceOf[CometShuffleManager] == expected)
+    withParquetTable((0 until 100).map(i => (i, (i % 7).toString)), "tbl") {
+      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 cometManager = classOf[CometShuffleManager].getName
+
+  private def chosen(conf: SparkConf): Option[String] = {
+    CometDriverPlugin.maybeSetShuffleManager(conf)
+    conf.getOption(SHUFFLE_MANAGER.key)
+  }
+
+  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(SHUFFLE_MANAGER.key, manager)).contains(manager))
+      }
+  }
+
+  test("leaves the shuffle manager unset when Comet or Comet shuffle is disabled") {
+    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 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 {