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