DynamicPartitionDataSingleWriter needs sort before write

Author: maomaodev

extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/write/HiveWrite.scala

@@ -32,7 +32,9 @@ import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.catalog.CatalogTable
 import org.apache.spark.sql.catalyst.catalog.CatalogTypes.TablePartitionSpec
-import org.apache.spark.sql.connector.write.{BatchWrite, LogicalWriteInfo, Write}
+import org.apache.spark.sql.connector.distributions.{Distribution, Distributions}
+import org.apache.spark.sql.connector.expressions.{Expressions, SortDirection, SortOrder}
+import org.apache.spark.sql.connector.write.{BatchWrite, LogicalWriteInfo, RequiresDistributionAndOrdering}
 import org.apache.spark.sql.execution.datasources.{BasicWriteJobStatsTracker, WriteJobDescription}
 import org.apache.spark.sql.execution.datasources.v2.FileBatchWrite
 import org.apache.spark.sql.execution.metric.SQLMetric
@@ -50,7 +52,8 @@ case class HiveWrite(
     info: LogicalWriteInfo,
     hiveTableCatalog: HiveTableCatalog,
     forceOverwrite: Boolean,
-    dynamicPartition: Map[String, Option[String]]) extends Write with Logging {
+    dynamicPartition: Map[String, Option[String]])
+  extends RequiresDistributionAndOrdering with Logging {
 
   private val options = info.options()
 
@@ -73,6 +76,14 @@ case class HiveWrite(
 
   override def description(): String = "Kyuubi-Hive-Connector"
 
+  override def requiredDistribution(): Distribution = Distributions.unspecified()
+
+  override def requiredOrdering(): Array[SortOrder] = {
+    partColumns.map { col =>
+      Expressions.sort(Expressions.column(col.name), SortDirection.ASCENDING)
+    }.toArray
+  }
+
   override def toBatch: BatchWrite = {
     val tmpLocation = HiveWriteHelper.getExternalTmpPath(externalCatalog, hadoopConf, tableLocation)
 

extensions/spark/kyuubi-spark-connector-hive/src/test/scala/org/apache/kyuubi/spark/connector/hive/HiveQuerySuite.scala

@@ -117,6 +117,55 @@ class HiveQuerySuite extends KyuubiHiveTest {
     }
   }
 
+  test("[KYUUBI #7335] DynamicPartitionDataSingleWriter needs sort before write") {
+    withSparkSession(Map(
+      "hive.exec.dynamic.partition.mode" -> "nonstrict",
+      "spark.sql.shuffle.partitions" -> "1")) { spark =>
+      val table = "hive.default.test_part_table"
+      val tempTable = "hive.default.test_part_table_tmp"
+      withTable(table, tempTable) {
+        spark.sql(
+          s"""
+             | CREATE TABLE $table (
+             |   word STRING,
+             |   num BIGINT
+             | ) PARTITIONED BY (dt STRING)
+             | STORED AS ORC
+             |""".stripMargin)
+        spark.sql(
+          s"""
+             | CREATE TABLE $tempTable (
+             |   word STRING,
+             |   num BIGINT,
+             |   dt STRING
+             | ) STORED AS ORC
+             |""".stripMargin)
+        spark.sql(
+          s"""
+             | INSERT INTO $tempTable VALUES
+             | ('1', 1, '1111'),
+             | ('2', 2, '2222'),
+             | ('3', 4, '1111')
+             |""".stripMargin)
+
+        spark.sql(
+          s"""
+             | INSERT OVERWRITE TABLE $table PARTITION (dt)
+             | SELECT word, num, dt
+             | FROM $tempTable
+             | ORDER BY word
+             |""".stripMargin).collect()
+
+        checkAnswer(
+          spark.sql(s"SELECT * FROM $table"),
+          Seq(
+            Row("1", 1L, "1111"),
+            Row("2", 2L, "2222"),
+            Row("3", 4L, "1111")))
+      }
+    }
+  }
+
   test("[KYUUBI #4525] Partitioning predicates should take effect to filter data") {
     withSparkSession(Map("hive.exec.dynamic.partition.mode" -> "nonstrict")) { spark =>
       val table = "hive.default.employee"