sunchao commented on code in PR #5957:
URL: https://github.com/apache/datafusion-comet/pull/5957#discussion_r4022334508
##########
spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala:
##########
@@ -19,18 +19,113 @@
package org.apache.comet.rules
-import org.apache.spark.sql.CometTestBase
+import org.apache.spark.sql.{CometTestBase, SaveMode}
+import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Literal}
import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial}
import org.apache.spark.sql.comet._
import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec
import org.apache.spark.sql.execution._
import org.apache.spark.sql.execution.adaptive.QueryStageExec
+import org.apache.spark.sql.execution.command.DataWritingCommandExec
+import org.apache.spark.sql.execution.datasources.WriteFilesExec
import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.BinaryType
+import org.apache.spark.sql.util.QueryExecutionListener
import org.apache.comet.CometConf
+import org.apache.comet.serde.OperatorOuterClass.Operator
+
+private case class AliasingFallbackCometExec(
+ override val originalPlan: SparkPlan,
+ child: SparkPlan)
+ extends CometExec
+ with UnaryExecNode {
+ override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
+ copy(child = newChild)
+}
class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase {
+ private def cometIcebergWrite(child: SparkPlan): CometIcebergWriteExec = {
+ val output = Seq(
+ AttributeReference(IcebergWriteExec.CommitMessageColumn, BinaryType,
nullable = false)())
+ val originalPlan = IcebergWriteExec(null, output, child)
+ CometIcebergWriteExec(
+ Operator.newBuilder().build(),
+ originalPlan,
+ child,
+ output,
+ batchWrite = null,
+ table = null,
+ partitionSpecId = 0)
+ }
+
+ private def cometFilter(child: SparkPlan): CometFilterExec = {
+ val condition = Literal.TrueLiteral
+ val sparkFilter = FilterExec(condition, child)
+ CometFilterExec(
+ Operator.newBuilder().build(),
+ sparkFilter,
+ sparkFilter.output,
+ condition,
+ child,
+ SerializedPlan(None))
+ }
+
+ private def captureDataWritingCommand(path: String): DataWritingCommandExec
= {
+ var captured: SparkPlan = null
+ val listener = new QueryExecutionListener {
+ override def onSuccess(funcName: String, qe: QueryExecution, durationNs:
Long): Unit = {
+ if (funcName == "save" || funcName.contains("command")) {
+ captured = qe.executedPlan
+ }
+ }
+ override def onFailure(
+ funcName: String,
+ qe: QueryExecution,
+ exception: Exception): Unit = {}
+ }
+ spark.listenerManager.register(listener)
+ try {
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ spark.range(1).toDF("id").write.mode("overwrite").parquet(path)
+ }
+ } finally {
+ spark.listenerManager.unregister(listener)
Review Comment:
### Correctness
[P2] Wait for the write callback before unregistering the listener
A successful `parquet(path)` call does not guarantee that
`QueryExecutionListener.onSuccess` has run: Spark posts the SQL execution-end
event to the shared asynchronous listener queue. If that queue delivers the
event after this `finally` block, the listener has already been removed and
`captured` stays null, so the three new Parquet reversion tests can fail with
`expected a captured parquet write plan` even though the write succeeded. I
verified this delivery path in the maintained Spark 3.5 and 4.0 branches and
reproduced the failing schedule with the exact helper in isolated Scala test
doubles. Immediate delivery and a wait-before-unregister control both pass.
Please wait with a bounded latch/future while the listener remains registered,
then unregister in `finally`, so completion and visibility of the captured plan
are guaranteed.
--
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]
