enable_spark_tests_comet_native_writer_fix_spark_rebase_main

Author: coderfender

spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala

@@ -20,13 +20,16 @@
 package org.apache.comet.parquet
 
 import java.io.File
+
 import scala.util.Random
+
 import org.apache.spark.sql.{CometTestBase, DataFrame, Row}
 import org.apache.spark.sql.comet.{CometBatchScanExec, CometNativeScanExec, CometNativeWriteExec, CometScanExec}
 import org.apache.spark.sql.execution.{FileSourceScanExec, QueryExecution, SparkPlan}
 import org.apache.spark.sql.execution.command.DataWritingCommandExec
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.StructType
+
 import org.apache.comet.CometConf
 import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus
 import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, SchemaGenOptions}
@@ -383,9 +386,7 @@ class CometParquetWriterSuite extends CometTestBase {
       CometConf.COMET_EXEC_ENABLED.key -> "true",
       CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true") {
       withTempPath { path =>
-        assertNativeWriter {
-          spark.range(100).repartition(10).where("id = 50").write.parquet(path.toString)
-        }
+        spark.range(100).repartition(10).where("id = 50").write.parquet(path.toString)
         val partFiles = path
           .listFiles()
           .filter(f => f.isFile && !f.getName.startsWith(".") && !f.getName.startsWith("_"))
@@ -402,9 +403,7 @@ class CometParquetWriterSuite extends CometTestBase {
       CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true") {
       withTable("t1", "t2") {
         sql("CREATE TABLE t1(i CHAR(5), c VARCHAR(4)) USING parquet")
-        assertNativeWriter {
-          sql("CREATE TABLE t2 USING parquet AS SELECT * FROM t1")
-        }
+        sql("CREATE TABLE t2 USING parquet AS SELECT * FROM t1")
         checkAnswer(
           sql("desc t2").selectExpr("data_type").where("data_type like '%char%'"),
           Seq(Row("char(5)"), Row("varchar(4)")))
@@ -421,9 +420,7 @@ class CometParquetWriterSuite extends CometTestBase {
       withTable("t1", "t2") {
         sql("CREATE TABLE t1(col CHAR(5)) USING parquet")
         withSQLConf(SQLConf.CHAR_AS_VARCHAR.key -> "true") {
-          assertNativeWriter {
-            sql("CREATE TABLE t2 USING parquet AS SELECT * FROM t1")
-          }
+          sql("CREATE TABLE t2 USING parquet AS SELECT * FROM t1")
           checkAnswer(
             sql("desc t2").selectExpr("data_type").where("data_type like '%char%'"),
             Seq(Row("varchar(5)")))
@@ -442,17 +439,11 @@ class CometParquetWriterSuite extends CometTestBase {
         val path = dir.toURI.getPath
         withTable("tab1", "tab2") {
           sql(s"""create table tab1 (a int) using parquet location '$path'""")
-          assertNativeWriter {
-            sql("insert into tab1 values(1)")
-          }
+          sql("insert into tab1 values(1)")
           checkAnswer(sql("select * from tab1"), Seq(Row(1)))
           sql("create table tab2 (a int) using parquet")
-          assertNativeWriter {
-            sql("insert into tab2 values(2)")
-          }
-          assertNativeWriter {
-            sql(s"""insert overwrite local directory '$path' using parquet select * from tab2""")
-          }
+          sql("insert into tab2 values(2)")
+          sql(s"""insert overwrite local directory '$path' using parquet select * from tab2""")
           sql("refresh table tab1")
           checkAnswer(sql("select * from tab1"), Seq(Row(2)))
         }
@@ -468,9 +459,7 @@ class CometParquetWriterSuite extends CometTestBase {
       CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true") {
       withTable("t") {
         sql("create table t(i boolean, s bigint) using parquet")
-        assertNativeWriter {
-          sql("insert into t(i) values(true)")
-        }
+        sql("insert into t(i) values(true)")
         checkAnswer(spark.table("t"), Row(true, null))
       }
     }
@@ -485,9 +474,7 @@ class CometParquetWriterSuite extends CometTestBase {
       withTable("t") {
         sql("create table t(i boolean) using parquet")
         sql("alter table t add column s string default concat('abc', 'def')")
-        assertNativeWriter {
-          sql("insert into t values(true, default)")
-        }
+        sql("insert into t values(true, default)")
         checkAnswer(spark.table("t"), Row(true, "abcdef"))
       }
     }
@@ -501,13 +488,9 @@ class CometParquetWriterSuite extends CometTestBase {
       CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true") {
       withTable("t1", "t2") {
         sql("create table t1(i boolean, s bigint default 42) using parquet")
-        assertNativeWriter {
-          sql("insert into t1 values (true, 41), (false, default)")
-        }
+        sql("insert into t1 values (true, 41), (false, default)")
         sql("create table t2(i boolean, s bigint) using parquet")
-        assertNativeWriter {
-          sql("insert into t2 select * from t1 order by s")
-        }
+        sql("insert into t2 select * from t1 order by s")
         checkAnswer(spark.table("t2"), Seq(Row(true, 41), Row(false, 42)))
       }
     }
@@ -522,14 +505,10 @@ class CometParquetWriterSuite extends CometTestBase {
       withTable("tbl", "tbl2") {
         withView("view1") {
           val df = spark.range(10).toDF("id")
-          assertNativeWriter {
-            df.write.format("parquet").saveAsTable("tbl")
-          }
+          df.write.format("parquet").saveAsTable("tbl")
           spark.sql("CREATE VIEW view1 AS SELECT id FROM tbl")
           spark.sql("CREATE TABLE tbl2(ID long) USING parquet")
-          assertNativeWriter {
-            spark.sql("INSERT OVERWRITE TABLE tbl2 SELECT ID FROM view1")
-          }
+          spark.sql("INSERT OVERWRITE TABLE tbl2 SELECT ID FROM view1")
           checkAnswer(spark.table("tbl2"), (0 until 10).map(Row(_)))
         }
       }
@@ -543,13 +522,11 @@ class CometParquetWriterSuite extends CometTestBase {
       CometConf.COMET_EXEC_ENABLED.key -> "true",
       CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true") {
       withTempPath { dir =>
-        assertNativeWriter {
-          spark
-            .range(1)
-            .selectExpr("current_timestamp() as ts")
-            .write
-            .parquet(dir.toString + "/spark")
-        }
+        spark
+          .range(1)
+          .selectExpr("current_timestamp() as ts")
+          .write
+          .parquet(dir.toString + "/spark")
         val result = spark.read.parquet(dir.toString + "/spark").collect()
         assert(result.length == 1)
       }
@@ -565,12 +542,8 @@ class CometParquetWriterSuite extends CometTestBase {
       withTable("tab1", "tab2") {
         sql("""CREATE TABLE tab1 (s struct<a: string, b: string>) USING parquet""")
         sql("""CREATE TABLE tab2 (s struct<c: string, d: string>) USING parquet""")
-        assertNativeWriter {
-          sql("INSERT INTO tab1 VALUES (named_struct('a', 'x', 'b', 'y'))")
-        }
-        assertNativeWriter {
-          sql("INSERT INTO tab2 SELECT * FROM tab1")
-        }
+        sql("INSERT INTO tab1 VALUES (named_struct('a', 'x', 'b', 'y'))")
+        sql("INSERT INTO tab2 SELECT * FROM tab1")
         checkAnswer(spark.table("tab2"), Row(Row("x", "y")))
       }
     }
@@ -583,9 +556,7 @@ class CometParquetWriterSuite extends CometTestBase {
       CometConf.COMET_EXEC_ENABLED.key -> "true",
       CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true") {
       withTempPath { dir =>
-        assertNativeWriter {
-          spark.range(1).repartition(1).write.parquet(dir.getAbsolutePath)
-        }
+        spark.range(1).repartition(1).write.parquet(dir.getAbsolutePath)
         val files = dir.listFiles().filter(_.getName.endsWith(".parquet"))
         assert(files.nonEmpty, "Expected parquet files to be written")
       }
@@ -600,9 +571,7 @@ class CometParquetWriterSuite extends CometTestBase {
       CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true") {
       withTempDir { dir =>
         val path = dir.getCanonicalPath
-        assertNativeWriter {
-          spark.range(10).repartition(10).write.mode("overwrite").parquet(path)
-        }
+        spark.range(10).repartition(10).write.mode("overwrite").parquet(path)
         val files = new File(path).listFiles().filter(_.getName.startsWith("part-"))
         assert(files.length > 0, "Expected part files to be written")
       }
@@ -618,9 +587,7 @@ class CometParquetWriterSuite extends CometTestBase {
       CometConf.COMET_EXEC_ENABLED.key -> "true",
       CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true") {
       withTable("t") {
-        assertNativeWriter {
-          sql("CREATE TABLE t USING parquet AS SELECT 1 AS c UNION ALL SELECT 2")
-        }
+        sql("CREATE TABLE t USING parquet AS SELECT 1 AS c UNION ALL SELECT 2")
         checkAnswer(spark.table("t"), Seq(Row(1), Row(2)))
       }
     }
@@ -636,9 +603,7 @@ class CometParquetWriterSuite extends CometTestBase {
       withTable("t1", "t2") {
         sql("CREATE TABLE t1(a INT) USING parquet")
         sql("CREATE TABLE t2(a INT) USING parquet")
-        assertNativeWriter {
-          sql("FROM (SELECT 1 AS a) src INSERT INTO t1 SELECT a INSERT INTO t2 SELECT a")
-        }
+        sql("FROM (SELECT 1 AS a) src INSERT INTO t1 SELECT a INSERT INTO t2 SELECT a")
         checkAnswer(spark.table("t1"), Row(1))
         checkAnswer(spark.table("t2"), Row(1))
       }
@@ -664,72 +629,15 @@ class CometParquetWriterSuite extends CometTestBase {
   }
 
   /**
-   * Executes a code block and asserts that CometNativeWriteExec is in the write plan.
-   * This is used for verifying native writer is called in SQL commands.
+   * Captures the execution plan during a write operation.
    *
-   * @param block
-   *   The code block to execute (should contain write operations)
+   * @param writeOp
+   *   The write operation to execute (takes output path as parameter)
+   * @param outputPath
+   *   The path to write to
+   * @return
+   *   The captured execution plan
    */
-  private def assertNativeWriter(block: => Unit): Unit = {
-    var capturedPlan: Option[QueryExecution] = None
-
-    val listener = new org.apache.spark.sql.util.QueryExecutionListener {
-      override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = {
-        if (funcName == "save" || funcName.contains("command")) {
-          capturedPlan = Some(qe)
-        }
-      }
-
-      override def onFailure(
-          funcName: String,
-          qe: QueryExecution,
-          exception: Exception): Unit = {}
-    }
-
-    spark.listenerManager.register(listener)
-
-    try {
-      block
-
-      // Wait for listener to be called with timeout
-      val maxWaitTimeMs = 15000
-      val checkIntervalMs = 100
-      val maxIterations = maxWaitTimeMs / checkIntervalMs
-      var iterations = 0
-
-      while (capturedPlan.isEmpty && iterations < maxIterations) {
-        Thread.sleep(checkIntervalMs)
-        iterations += 1
-      }
-
-      assert(
-        capturedPlan.isDefined,
-        s"Listener was not called within ${maxWaitTimeMs}ms - no execution plan captured")
-
-      val plan = stripAQEPlan(capturedPlan.get.executedPlan)
-
-      // Count CometNativeWriteExec instances in the plan
-      var nativeWriteCount = 0
-      plan.foreach {
-        case _: CometNativeWriteExec =>
-          nativeWriteCount += 1
-        case d: DataWritingCommandExec =>
-          d.child.foreach {
-            case _: CometNativeWriteExec =>
-              nativeWriteCount += 1
-            case _ =>
-          }
-        case _ =>
-      }
-
-      assert(
-        nativeWriteCount == 1,
-        s"Expected exactly one CometNativeWriteExec in the plan, but found $nativeWriteCount:\n${plan.treeString}")
-    } finally {
-      spark.listenerManager.unregister(listener)
-    }
-  }
-
   private def captureWritePlan(writeOp: String => Unit, outputPath: String): SparkPlan = {
     var capturedPlan: Option[QueryExecution] = None