add_listner_comet_plugin_ini

Author: coderfender

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

@@ -19,22 +19,11 @@
 
 package org.apache.comet
 
-import java.util.concurrent.atomic.AtomicBoolean
-
 import org.apache.spark.CometSource
-import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.execution.QueryExecution
 import org.apache.spark.sql.util.QueryExecutionListener
 
-object CometMetricsListener extends QueryExecutionListener {
-
-  private val registered = new AtomicBoolean(false)
-
-  def ensureRegistered(session: SparkSession): Unit = {
-    if (registered.compareAndSet(false, true)) {
-      session.listenerManager.register(this)
-    }
-  }
+class CometMetricsListener extends QueryExecutionListener {
 
   override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = {
     val stats = CometCoverageStats.forPlan(qe.executedPlan)

spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala

@@ -47,7 +47,7 @@ import org.apache.spark.sql.execution.window.WindowExec
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
 
-import org.apache.comet.{CometConf, CometExplainInfo, CometMetricsListener, ExtendedExplainInfo}
+import org.apache.comet.{CometConf, CometExplainInfo, ExtendedExplainInfo}
 import org.apache.comet.CometConf.{COMET_SPARK_TO_ARROW_ENABLED, COMET_SPARK_TO_ARROW_SUPPORTED_OPERATOR_LIST}
 import org.apache.comet.CometSparkSessionExtensions._
 import org.apache.comet.rules.CometExecRule.allExecs
@@ -387,8 +387,6 @@ case class CometExecRule(session: SparkSession) extends Rule[SparkPlan] {
         normalizedPlan
       }
 
-      CometMetricsListener.ensureRegistered(session)
-
       var newPlan = transform(planWithJoinRewritten)
 
       // if the plan cannot be run fully natively then explain why (when appropriate

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

@@ -60,6 +60,9 @@ class CometDriverPlugin extends DriverPlugin with Logging with ShimCometDriverPl
     // Register Comet metrics source
     sc.env.metricsSystem.registerSource(CometSource)
 
+    // Register query execution listener via config
+    CometDriverPlugin.registerQueryExecutionListener(sc.conf)
+
     if (CometSparkSessionExtensions.shouldOverrideMemoryConf(sc.getConf)) {
       val execMemOverhead = if (sc.getConf.contains(EXECUTOR_MEMORY_OVERHEAD.key)) {
         sc.getConf.getSizeAsMb(EXECUTOR_MEMORY_OVERHEAD.key)
@@ -104,6 +107,10 @@ class CometDriverPlugin extends DriverPlugin with Logging with ShimCometDriverPl
 }
 
 object CometDriverPlugin extends Logging {
+  def registerQueryExecutionListener(conf: SparkConf): Unit = {
+    conf.set("spark.sql.queryExecutionListeners", "org.apache.comet.CometMetricsListener")
+  }
+
   def registerCometSessionExtension(conf: SparkConf): Unit = {
     val extensionKey = StaticSQLConf.SPARK_SESSION_EXTENSIONS.key
     val extensionClass = classOf[CometSparkSessionExtensions].getName

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

@@ -88,7 +88,7 @@ class CometPluginsSuite extends CometTestBase {
       spark.range(1000).toDF("id").write.mode(SaveMode.Overwrite).parquet(path)
       spark.read.parquet(path).filter("id > 500").collect()
     }
-
+    spark.sparkContext.listenerBus.waitUntilEmpty()
     assert(
       CometSource.QUERIES_PLANNED.getCount > queriesBefore,
       "queries.planned should increment after query")
@@ -104,10 +104,10 @@ class CometPluginsSuite extends CometTestBase {
         spark.range(10000).toDF("id").write.mode(SaveMode.Overwrite).parquet(path)
 
         val queriesBefore = CometSource.QUERIES_PLANNED.getCount
-
+        spark.sparkContext.listenerBus.waitUntilEmpty()
         spark.read.parquet(path).filter("id > 100").collect()
         spark.read.parquet(path).filter("id > 200").collect()
-
+        spark.sparkContext.listenerBus.waitUntilEmpty()
         val queriesAfter = CometSource.QUERIES_PLANNED.getCount
         assert(
           queriesAfter == queriesBefore + 2,