sam-1112 commented on code in PR #5957:
URL: https://github.com/apache/datafusion-comet/pull/5957#discussion_r4022531054
##########
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:
Fixed. The helper now waits on a bounded `CountDownLatch` from `onSuccess`
while the listener is still registered, then unregisters in `finally`.
--
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]
