fix: warn about the deprecated memoryPool.fraction when it is set on the session

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

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

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 " +

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