wForget commented on code in PR #3708:
URL: https://github.com/apache/datafusion-comet/pull/3708#discussion_r2938636399
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -389,6 +390,10 @@ case class CometExecRule(session: SparkSession) extends
Rule[SparkPlan] {
var newPlan = transform(planWithJoinRewritten)
+ // Record coverage stats for metrics
+ val stats = CometCoverageStats.forPlan(newPlan)
Review Comment:
I mean that recording stats here is incorrect, because CometExecRule may be
executed multiple times within a single sql execution. So we might need to use
a QueryExecutionListener, like:
```
val queryExecutionListener = new QueryExecutionListener {
override def onSuccess(funcName: String, qe: QueryExecution,
durationNs: Long): Unit = {
recordStats(qe)
}
override def onFailure(funcName: String, qe: QueryExecution,
exception: Exception): Unit = {
recordStats(qe)
}
private def recordStats(qe: QueryExecution): Unit = {
try {
val plan = qe.executedPlan
val stats = CometCoverageStats.forPlan(plan)
CometSource.recordStats(stats)
} catch {
case _: Throwable => // ignore any exceptions to avoid test
flakiness
}
}
}
spark.listenerManager.register(queryExecutionListener)
```
The test case like:
```
val nativeBefore = CometSource.NATIVE_OPERATORS.getCount
val queriesBefore = CometSource.QUERIES_PLANNED.getCount
import testImplicits._
spark.read.parquet(path).repartition($"id").filter("id >
500").collect()
spark.sparkContext.listenerBus.waitUntilEmpty()
val nativeAfter= CometSource.NATIVE_OPERATORS.getCount
val queriesAfter = CometSource.QUERIES_PLANNED.getCount
assert(queriesAfter - queriesBefore == 1)
assert(nativeAfter - nativeBefore == 2)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
