From 7c87af8122f785bfcebb09eba23559a2b5a8bb27 Mon Sep 17 00:00:00 2001 From: tlfua <0050211@gmail.com> Date: Sat, 7 Nov 2026 16:58:35 -0800 Subject: [PATCH] fix: warn about the deprecated memoryPool.fraction when it is set on the session The deprecation warning for spark.comet.exec.memoryPool.fraction only checked the SparkConf in the driver plugin, but the value is read from the session conf for each query. A value set with SET or spark.conf.set, or in an application that registers the session extension without the plugin, got no warning. CometExecRule now checks the session conf too, and the warning is logged once per JVM. Also state in the memory tuning guide that fair_unified limits each task to the fraction times spark.memory.offHeap.size. Part of https://github.com/apache/datafusion-comet/issues/6187 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 1da6c065-9ae0-4ad5-a84d-12fe9e45d920 --- .../source/user-guide/latest/tuning/memory.md | 5 +-- .../apache/comet/rules/CometExecRule.scala | 4 +++ .../main/scala/org/apache/spark/Plugins.scala | 14 +++++--- .../org/apache/spark/CometPluginsSuite.scala | 34 ++++++++++++++----- 4 files changed, 42 insertions(+), 15 deletions(-) diff --git a/docs/source/user-guide/latest/tuning/memory.md b/docs/source/user-guide/latest/tuning/memory.md index 25510c8e9c3..d1d7b8ed75a 100644 --- a/docs/source/user-guide/latest/tuning/memory.md +++ b/docs/source/user-guide/latest/tuning/memory.md @@ -61,8 +61,9 @@ while Comet runs; see [Sizing the Overhead from the Memory Usage Log]. `spark.comet.exec.memoryPool.fraction` is deprecated and does not leave room for it. Spark hands out all of `spark.memory.offHeap.size` to the tasks that ask for it, whatever the fraction. The `fair_unified` pool applies the -fraction to each task separately, where Spark's own limit of an even share of the pool per running task is tighter -unless the fraction is below one over the number of running tasks, and the `greedy_unified` pool ignores it. +fraction to each task separately, limiting each task to the fraction times `spark.memory.offHeap.size`, where Spark's +own limit of an even share of the pool per running task is tighter unless the fraction is below one over the number +of running tasks, and the `greedy_unified` pool ignores it. For more details about Spark off-heap memory mode, please refer to [Spark documentation]. diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index 70d108f2506..dfff435e27a 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -24,6 +24,7 @@ import java.util.concurrent.ConcurrentHashMap import scala.collection.mutable.ListBuffer import scala.jdk.CollectionConverters._ +import org.apache.spark.CometDriverPlugin import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, Expression, SortOrder} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, Final, Partial, PartialMerge} @@ -1105,6 +1106,9 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) return plan } + CometDriverPlugin.warnIfMemoryPoolFractionSet( + Option(conf.getConfString(CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key, null))) + // Comet does not support structured streaming. Fall back to Spark for any plan that // belongs to a streaming query (detected via StreamSourceAwareSparkPlan.getStream). if (ShimCometStreaming.isStreamingPlan(plan)) return plan diff --git a/spark/src/main/scala/org/apache/spark/Plugins.scala b/spark/src/main/scala/org/apache/spark/Plugins.scala index 70f3ce6cedc..9f19577a763 100644 --- a/spark/src/main/scala/org/apache/spark/Plugins.scala +++ b/spark/src/main/scala/org/apache/spark/Plugins.scala @@ -21,7 +21,7 @@ package org.apache.spark import java.{util => ju} import java.util.Collections -import java.util.concurrent.atomic.AtomicReference +import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference} import scala.collection.mutable import scala.util.Try @@ -123,7 +123,8 @@ class CometDriverPlugin private[spark] (clock: Clock) extends DriverPlugin with CometDriverPlugin.registerIcebergWriteReport(sc.conf) CometDriverPlugin.warnIfExecutorMemoryOverheadUnset(sc.getConf) - CometDriverPlugin.warnIfMemoryPoolFractionSet(sc.getConf) + CometDriverPlugin.warnIfMemoryPoolFractionSet( + sc.getConf.getOption(CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key)) extraConfs } @@ -324,9 +325,14 @@ object CometDriverPlugin extends Logging { // for the native memory that Comet does not reserve. It cannot: Spark hands out the whole pool // to the tasks that ask for it, and the fraction only caps each task's total under // fair_unified. Users who set it for that purpose need to size the memory overhead instead. - private[apache] def warnIfMemoryPoolFractionSet(conf: SparkConf): Unit = { + // The driver plugin checks the SparkConf, and CometExecRule checks the session conf of each + // query, which also covers SET, spark.conf.set and sessions that register the extension + // without the plugin. The warning is logged once per JVM. + private[apache] val memoryPoolFractionWarned = new AtomicBoolean(false) + + private[apache] def warnIfMemoryPoolFractionSet(setTo: Option[String]): Unit = { val key = CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key - conf.getOption(key).foreach { value => + setTo.filter(_ => memoryPoolFractionWarned.compareAndSet(false, true)).foreach { value => logWarning( s"$key=$value is deprecated and will be removed in a future major release. It does " + "not leave room in spark.memory.offHeap.size for native memory that Comet's memory " + diff --git a/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala b/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala index f85045a57ec..da083e437d9 100644 --- a/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala @@ -278,25 +278,41 @@ class CometPluginsMemoryOverheadWarningSuite extends CometTestBase { class CometPluginsMemoryPoolFractionWarningSuite extends CometTestBase { - private val warning = "spark.comet.exec.memoryPool.fraction=0.8 is deprecated" + private val key = CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key + private val warning = s"$key=0.8 is deprecated" - private def warningsFor(conf: SparkConf): Seq[String] = { + private def warningsDuring(f: => Unit): Seq[String] = { // Logging derives the logger name by stripping the object's trailing '$' val logger = CometDriverPlugin.getClass.getName.stripSuffix("$") val appender = new LogAppender("memory pool fraction warning") - withLogAppender(appender, Seq(logger), Some(Level.WARN)) { - CometDriverPlugin.warnIfMemoryPoolFractionSet(conf) - } + CometDriverPlugin.memoryPoolFractionWarned.set(false) + withLogAppender(appender, Seq(logger), Some(Level.WARN))(f) appender.loggingEvents.map(_.getMessage.getFormattedMessage).toSeq } - test("warns when the memory pool fraction is set") { - val conf = new SparkConf().set("spark.comet.exec.memoryPool.fraction", "0.8") - assert(warningsFor(conf).exists(_.contains(warning))) + test("warns when the memory pool fraction is set in the SparkConf") { + val conf = new SparkConf().set(key, "0.8") + val warnings = + warningsDuring(CometDriverPlugin.warnIfMemoryPoolFractionSet(conf.getOption(key))) + assert(warnings.exists(_.contains(warning))) } test("does not warn when the memory pool fraction is unset") { - assert(!warningsFor(new SparkConf()).exists(_.contains("memoryPool.fraction"))) + val warnings = warningsDuring { + CometDriverPlugin.warnIfMemoryPoolFractionSet(new SparkConf().getOption(key)) + sql("SELECT 1").collect() + } + assert(!warnings.exists(_.contains("memoryPool.fraction"))) + } + + test("warns once when the memory pool fraction is set on the session") { + val warnings = warningsDuring { + withSQLConf(key -> "0.8") { + sql("SELECT 1").collect() + sql("SELECT 2").collect() + } + } + assert(warnings.count(_.contains(warning)) == 1) } }