enable_spark_tests_comet_native_writer_fix_spark_rebase_main

Author: coderfender

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

@@ -383,7 +383,9 @@ class CometParquetWriterSuite extends CometTestBase {
       CometConf.COMET_EXEC_ENABLED.key -> "true",
       CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true") {
       withTempPath { path =>
-        spark.range(100).repartition(10).where("id = 50").write.parquet(path.toString)
+        assertNativeWriter {
+          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("_"))
@@ -400,7 +402,9 @@ class CometParquetWriterSuite extends CometTestBase {
       CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true") {
       withTable("t1", "t2") {
         sql("CREATE TABLE t1(i CHAR(5), c VARCHAR(4)) USING parquet")
-        sql("CREATE TABLE t2 USING parquet AS SELECT * FROM t1")
+        assertNativeWriter {
+          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)")))
@@ -417,7 +421,9 @@ class CometParquetWriterSuite extends CometTestBase {
       withTable("t1", "t2") {
         sql("CREATE TABLE t1(col CHAR(5)) USING parquet")
         withSQLConf(SQLConf.CHAR_AS_VARCHAR.key -> "true") {
-          sql("CREATE TABLE t2 USING parquet AS SELECT * FROM t1")
+          assertNativeWriter {
+            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)")))
@@ -436,11 +442,17 @@ class CometParquetWriterSuite extends CometTestBase {
         val path = dir.toURI.getPath
         withTable("tab1", "tab2") {
           sql(s"""create table tab1 (a int) using parquet location '$path'""")
-          sql("insert into tab1 values(1)")
+          assertNativeWriter {
+            sql("insert into tab1 values(1)")
+          }
           checkAnswer(sql("select * from tab1"), Seq(Row(1)))
           sql("create table tab2 (a int) using parquet")
-          sql("insert into tab2 values(2)")
-          sql(s"""insert overwrite local directory '$path' using parquet select * from tab2""")
+          assertNativeWriter {
+            sql("insert into tab2 values(2)")
+          }
+          assertNativeWriter {
+            sql(s"""insert overwrite local directory '$path' using parquet select * from tab2""")
+          }
           sql("refresh table tab1")
           checkAnswer(sql("select * from tab1"), Seq(Row(2)))
         }
@@ -456,7 +468,9 @@ class CometParquetWriterSuite extends CometTestBase {
       CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true") {
       withTable("t") {
         sql("create table t(i boolean, s bigint) using parquet")
-        sql("insert into t(i) values(true)")
+        assertNativeWriter {
+          sql("insert into t(i) values(true)")
+        }
         checkAnswer(spark.table("t"), Row(true, null))
       }
     }
@@ -471,7 +485,9 @@ 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')")
-        sql("insert into t values(true, default)")
+        assertNativeWriter {
+          sql("insert into t values(true, default)")
+        }
         checkAnswer(spark.table("t"), Row(true, "abcdef"))
       }
     }
@@ -485,9 +501,13 @@ class CometParquetWriterSuite extends CometTestBase {
       CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true") {
       withTable("t1", "t2") {
         sql("create table t1(i boolean, s bigint default 42) using parquet")
-        sql("insert into t1 values (true, 41), (false, default)")
+        assertNativeWriter {
+          sql("insert into t1 values (true, 41), (false, default)")
+        }
         sql("create table t2(i boolean, s bigint) using parquet")
-        sql("insert into t2 select * from t1 order by s")
+        assertNativeWriter {
+          sql("insert into t2 select * from t1 order by s")
+        }
         checkAnswer(spark.table("t2"), Seq(Row(true, 41), Row(false, 42)))
       }
     }
@@ -502,10 +522,14 @@ class CometParquetWriterSuite extends CometTestBase {
       withTable("tbl", "tbl2") {
         withView("view1") {
           val df = spark.range(10).toDF("id")
-          df.write.format("parquet").saveAsTable("tbl")
+          assertNativeWriter {
+            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")
-          spark.sql("INSERT OVERWRITE TABLE tbl2 SELECT ID FROM view1")
+          assertNativeWriter {
+            spark.sql("INSERT OVERWRITE TABLE tbl2 SELECT ID FROM view1")
+          }
           checkAnswer(spark.table("tbl2"), (0 until 10).map(Row(_)))
         }
       }
@@ -519,11 +543,13 @@ class CometParquetWriterSuite extends CometTestBase {
       CometConf.COMET_EXEC_ENABLED.key -> "true",
       CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true") {
       withTempPath { dir =>
-        spark
-          .range(1)
-          .selectExpr("current_timestamp() as ts")
-          .write
-          .parquet(dir.toString + "/spark")
+        assertNativeWriter {
+          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)
       }
@@ -539,8 +565,12 @@ 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""")
-        sql("INSERT INTO tab1 VALUES (named_struct('a', 'x', 'b', 'y'))")
-        sql("INSERT INTO tab2 SELECT * FROM tab1")
+        assertNativeWriter {
+          sql("INSERT INTO tab1 VALUES (named_struct('a', 'x', 'b', 'y'))")
+        }
+        assertNativeWriter {
+          sql("INSERT INTO tab2 SELECT * FROM tab1")
+        }
         checkAnswer(spark.table("tab2"), Row(Row("x", "y")))
       }
     }
@@ -553,7 +583,9 @@ class CometParquetWriterSuite extends CometTestBase {
       CometConf.COMET_EXEC_ENABLED.key -> "true",
       CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true") {
       withTempPath { dir =>
-        spark.range(1).repartition(1).write.parquet(dir.getAbsolutePath)
+        assertNativeWriter {
+          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")
       }
@@ -568,7 +600,9 @@ class CometParquetWriterSuite extends CometTestBase {
       CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true") {
       withTempDir { dir =>
         val path = dir.getCanonicalPath
-        spark.range(10).repartition(10).write.mode("overwrite").parquet(path)
+        assertNativeWriter {
+          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")
       }
@@ -584,7 +618,9 @@ class CometParquetWriterSuite extends CometTestBase {
       CometConf.COMET_EXEC_ENABLED.key -> "true",
       CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true") {
       withTable("t") {
-        sql("CREATE TABLE t USING parquet AS SELECT 1 AS c UNION ALL SELECT 2")
+        assertNativeWriter {
+          sql("CREATE TABLE t USING parquet AS SELECT 1 AS c UNION ALL SELECT 2")
+        }
         checkAnswer(spark.table("t"), Seq(Row(1), Row(2)))
       }
     }
@@ -600,7 +636,9 @@ class CometParquetWriterSuite extends CometTestBase {
       withTable("t1", "t2") {
         sql("CREATE TABLE t1(a INT) USING parquet")
         sql("CREATE TABLE t2(a INT) USING parquet")
-        sql("FROM (SELECT 1 AS a) src INSERT INTO t1 SELECT a INSERT INTO t2 SELECT a")
+        assertNativeWriter {
+          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))
       }
@@ -626,15 +664,72 @@ class CometParquetWriterSuite extends CometTestBase {
   }
 
   /**
-   * Captures the execution plan during a write operation.
+   * 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.
    *
-   * @param writeOp
-   *   The write operation to execute (takes output path as parameter)
-   * @param outputPath
-   *   The path to write to
-   * @return
-   *   The captured execution plan
+   * @param block
+   *   The code block to execute (should contain write operations)
    */
+  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