Support save WithCTE for insertRepartitionBeforeWrite

extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/RepartitionBeforeWritingBase.scala

@@ -128,6 +128,7 @@ trait RepartitionBeforeWriteHelper extends Rule[LogicalPlan] {
   def canInsertRepartitionByExpression(plan: LogicalPlan): Boolean = {
     def canInsert(p: LogicalPlan): Boolean = p match {
       case Project(_, child) => canInsert(child)
+      case WithCTE(plan, _) => canInsert(plan)
       case SubqueryAlias(_, child) => canInsert(child)
       case Limit(_, _) => false
       case _: Sort => false

extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/kyuubi/sql/RepartitionBeforeWritingBase.scala

@@ -109,6 +109,7 @@ trait RepartitionBeforeWriteHelper extends Rule[LogicalPlan] {
   def canInsertRepartitionByExpression(plan: LogicalPlan): Boolean = {
     def canInsert(p: LogicalPlan): Boolean = p match {
       case Project(_, child) => canInsert(child)
+      case WithCTE(plan, _) => canInsert(plan)
       case SubqueryAlias(_, child) => canInsert(child)
       case Limit(_, _) => false
       case _: Sort => false

extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/RepartitionBeforeWritingBase.scala

@@ -109,6 +109,7 @@ trait RepartitionBeforeWriteHelper extends Rule[LogicalPlan] {
   def canInsertRepartitionByExpression(plan: LogicalPlan): Boolean = {
     def canInsert(p: LogicalPlan): Boolean = p match {
       case Project(_, child) => canInsert(child)
+      case WithCTE(plan, _) => canInsert(plan)
       case SubqueryAlias(_, child) => canInsert(child)
       case Limit(_, _) => false
       case _: Sort => false

extensions/spark/kyuubi-extension-spark-3-5/src/test/scala/org/apache/spark/sql/KyuubiSparkSQLExtensionTest.scala

@@ -100,6 +100,10 @@ trait KyuubiSparkSQLExtensionTest extends QueryTest
   }
 
   def withListener(df: => DataFrame)(callback: DataWritingCommand => Unit): Unit = {
+    runWithListener(df.collect())(callback)
+  }
+
+  def runWithListener(func: => Unit)(callback: DataWritingCommand => Unit): Unit = {
     val listener = new QueryExecutionListener {
       override def onFailure(f: String, qe: QueryExecution, e: Exception): Unit = {}
 
@@ -112,7 +116,7 @@ trait KyuubiSparkSQLExtensionTest extends QueryTest
     }
     spark.listenerManager.register(listener)
     try {
-      df.collect()
+      func
       sparkContext.listenerBus.waitUntilEmpty()
     } finally {
       spark.listenerManager.unregister(listener)

extensions/spark/kyuubi-extension-spark-3-5/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala

@@ -292,4 +292,36 @@ class RebalanceBeforeWritingSuite extends KyuubiSparkSQLExtensionTest {
       }
     }
   }
+
+  test("check rebalance exists for WithCTE") {
+    withSQLConf(
+      KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE.key -> "true",
+      KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE_IF_NO_SHUFFLE.key -> "true") {
+      withTempDir { tempDir =>
+        withTable("tmp1") {
+          spark.sql("CREATE TABLE tmp1 (id BIGINT,name BIGINT)")
+
+          val result = sql(
+            """
+              |WITH tmp_table AS (
+              |  SELECT id, name FROM tmp1 GROUP BY id, name
+              |)
+              |SELECT id, name FROM tmp_table ORDER BY id""".stripMargin)
+
+          val colName = (0 until result.schema.size).map(x => "col" + x)
+          val renamedDf = result.toDF(colName: _*)
+
+          runWithListener {
+            renamedDf.write
+              .mode("overwrite")
+              .save(tempDir.getCanonicalPath)
+          } { write =>
+            assert(write.collect {
+              case r: RebalancePartitions => r
+            }.size == 1)
+          }
+        }
+      }
+    }
+  }
 }