[SPARK-29037] For static partition overwrite, spark may give duplicate result.

Author: turboFei

core/src/main/scala/org/apache/spark/internal/io/FileCommitProtocol.scala

@@ -147,7 +147,9 @@ object FileCommitProtocol extends Logging {
       className: String,
       jobId: String,
       outputPath: String,
-      dynamicPartitionOverwrite: Boolean = false): FileCommitProtocol = {
+      dynamicPartitionOverwrite: Boolean = false,
+      staticPartitionKVS: Seq[(String, String)] = Seq.empty[(String, String)]):
+  FileCommitProtocol = {
 
     logDebug(s"Creating committer $className; job $jobId; output=$outputPath;" +
       s" dynamic=$dynamicPartitionOverwrite")
@@ -156,9 +158,11 @@ object FileCommitProtocol extends Logging {
     // dynamicPartitionOverwrite: Boolean).
     // If that doesn't exist, try the one with (jobId: string, outputPath: String).
     try {
-      val ctor = clazz.getDeclaredConstructor(classOf[String], classOf[String], classOf[Boolean])
-      logDebug("Using (String, String, Boolean) constructor")
-      ctor.newInstance(jobId, outputPath, dynamicPartitionOverwrite.asInstanceOf[java.lang.Boolean])
+      val ctor = clazz.getDeclaredConstructor(classOf[String], classOf[String], classOf[Boolean],
+        classOf[Seq[(String, String)]])
+      logDebug("Using (String, String, Boolean, Seq[(String, String)]) constructor")
+      ctor.newInstance(jobId, outputPath, dynamicPartitionOverwrite.asInstanceOf[java.lang.Boolean],
+        staticPartitionKVS)
     } catch {
       case _: NoSuchMethodException =>
         logDebug("Falling back to (String, String) constructor")

core/src/main/scala/org/apache/spark/internal/io/HadoopMapReduceCommitProtocol.scala

@@ -17,7 +17,7 @@
 
 package org.apache.spark.internal.io
 
-import java.io.IOException
+import java.io.{File, IOException}
 import java.util.{Date, UUID}
 
 import scala.collection.mutable
@@ -26,7 +26,7 @@ import scala.util.Try
 import org.apache.hadoop.conf.Configurable
 import org.apache.hadoop.fs.Path
 import org.apache.hadoop.mapreduce._
-import org.apache.hadoop.mapreduce.lib.output.FileOutputCommitter
+import org.apache.hadoop.mapreduce.lib.output.{FileOutputCommitter, FileOutputFormat}
 import org.apache.hadoop.mapreduce.task.TaskAttemptContextImpl
 
 import org.apache.spark.internal.Logging
@@ -52,7 +52,8 @@ import org.apache.spark.mapred.SparkHadoopMapRedUtil
 class HadoopMapReduceCommitProtocol(
     jobId: String,
     path: String,
-    dynamicPartitionOverwrite: Boolean = false)
+    dynamicPartitionOverwrite: Boolean = false,
+    staticPartitionKVS: Seq[(String, String)] = Seq.empty[(String, String)])
   extends FileCommitProtocol with Serializable with Logging {
 
   import FileCommitProtocol._
@@ -89,9 +90,15 @@ class HadoopMapReduceCommitProtocol(
    * The staging directory of this write job. Spark uses it to deal with files with absolute output
    * path, or writing data into partitioned directory with dynamicPartitionOverwrite=true.
    */
-  private def stagingDir = new Path(path, ".spark-staging-" + jobId)
+  protected def stagingDir = new Path(path, ".spark-staging-" + jobId)
+
+
+  private def getStaticPartitionPath(): String = {
+    staticPartitionKVS.map(kv => kv._1 + "=" + kv._2).mkString(File.separator)
+  }
 
   protected def setupCommitter(context: TaskAttemptContext): OutputCommitter = {
+    context.getConfiguration.set(FileOutputFormat.OUTDIR, stagingDir.toString)
     val format = context.getOutputFormatClass.getConstructor().newInstance()
     // If OutputFormat is Configurable, we should set conf to it.
     format match {
@@ -200,6 +207,20 @@ class HadoopMapReduceCommitProtocol(
           }
           fs.rename(new Path(stagingDir, part), finalPartPath)
         }
+      } else if (!getStaticPartitionPath().isEmpty) {
+        val finalPartPath = new Path(path, getStaticPartitionPath)
+        assert(!fs.exists(finalPartPath))
+        fs.rename(new Path(stagingDir, getStaticPartitionPath), finalPartPath)
+      } else {
+        val parts = fs.listStatus(stagingDir)
+          .filter(_.isDirectory)
+          .map(_.getPath.getName)
+          .filter(name => !name.startsWith(".") && name.contains("="))
+        for (part <- parts) {
+          val finalPartPath = new Path(path, part)
+          assert(!fs.exists(finalPartPath))
+          fs.rename(new Path(stagingDir, part), finalPartPath)
+        }
       }
 
       fs.delete(stagingDir, true)

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/InsertIntoHadoopFsRelationCommand.scala

@@ -105,11 +105,16 @@ case class InsertIntoHadoopFsRelationCommand(
     val dynamicPartitionOverwrite = enableDynamicOverwrite && mode == SaveMode.Overwrite &&
       staticPartitions.size < partitionColumns.length
 
+    val staticPartitionKVs = partitionColumns
+      .filter(c => staticPartitions.contains(c.name))
+      .map(att => (att.name, staticPartitions.get(att.name).get))
+
     val committer = FileCommitProtocol.instantiate(
       sparkSession.sessionState.conf.fileCommitProtocolClass,
-      jobId = java.util.UUID.randomUUID().toString,
-      outputPath = outputPath.toString,
-      dynamicPartitionOverwrite = dynamicPartitionOverwrite)
+      java.util.UUID.randomUUID().toString,
+      outputPath.toString,
+      dynamicPartitionOverwrite,
+      staticPartitionKVs)
 
     val doInsertion = (mode, pathExists) match {
       case (SaveMode.ErrorIfExists, true) =>

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SQLHadoopMapReduceCommitProtocol.scala

@@ -32,8 +32,9 @@ import org.apache.spark.sql.internal.SQLConf
 class SQLHadoopMapReduceCommitProtocol(
     jobId: String,
     path: String,
-    dynamicPartitionOverwrite: Boolean = false)
-  extends HadoopMapReduceCommitProtocol(jobId, path, dynamicPartitionOverwrite)
+    dynamicPartitionOverwrite: Boolean = false,
+    staticPartitionKVS: Seq[(String, String)] = Seq.empty[(String, String)])
+  extends HadoopMapReduceCommitProtocol(jobId, path, dynamicPartitionOverwrite, staticPartitionKVS)
     with Serializable with Logging {
 
   override protected def setupCommitter(context: TaskAttemptContext): OutputCommitter = {
@@ -55,7 +56,7 @@ class SQLHadoopMapReduceCommitProtocol(
         // The specified output committer is a FileOutputCommitter.
         // So, we will use the FileOutputCommitter-specified constructor.
         val ctor = clazz.getDeclaredConstructor(classOf[Path], classOf[TaskAttemptContext])
-        committer = ctor.newInstance(new Path(path), context)
+        committer = ctor.newInstance(stagingDir, context)
       } else {
         // The specified output committer is just an OutputCommitter.
         // So, we will use the no-argument constructor.

sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala

@@ -22,6 +22,7 @@ import java.net.{MalformedURLException, URL}
 import java.sql.{Date, Timestamp}
 import java.util.concurrent.atomic.AtomicBoolean
 
+import com.google.common.io.Files
 import org.apache.spark.{AccumulatorSuite, SparkException}
 import org.apache.spark.scheduler.{SparkListener, SparkListenerJobStart}
 import org.apache.spark.sql.catalyst.util.StringUtils
@@ -33,6 +34,7 @@ import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan
 import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, CartesianProductExec, SortMergeJoinExec}
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.internal.SQLConf.{PARTITION_OVERWRITE_MODE, PartitionOverwriteMode}
 import org.apache.spark.sql.test.{SharedSparkSession, TestSQLContext}
 import org.apache.spark.sql.test.SQLTestData._
 import org.apache.spark.sql.types._
@@ -3192,6 +3194,29 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession {
       checkAnswer(df3, Array(Row(new java.math.BigDecimal("0.100000000000000000000000100"))))
     }
   }
+
+  test("SPARK-29037: For non dynamic partition overwrite, set a unique staging dir") {
+    withSQLConf(PARTITION_OVERWRITE_MODE.key -> PartitionOverwriteMode.STATIC.toString) {
+      withTable("test") {
+        sql("create table test(id int, p1 int, p2 int) using parquet partitioned by (p1, p2)")
+        sql("insert overwrite table test partition(p1=1,p2) select 1, 3")
+        val df1 = sql("select * from test order by p2")
+        checkAnswer(df1, Array(Row(1, 1, 3)))
+        sql("insert overwrite table test partition(p1=1,p2) select 1, 4")
+        val df2 = sql("select * from test order by p2")
+        checkAnswer(df2, Array(Row(1, 1, 4)))
+        sql("insert overwrite table test partition(p1=1,p2=5) select 1")
+        val df3 = sql("select * from test order by p2")
+        checkAnswer(df3, Array(Row(1, 1, 4), Row(1, 1, 5)))
+        sql("insert overwrite table test select 1, 2, 3")
+        val df4 = sql("select * from test order by p2")
+        checkAnswer(df4, Array(Row(1, 2, 3)))
+        sql("insert overwrite table test select 9, 9, 9")
+        val df5 = sql("select * from test order by p2")
+        checkAnswer(df5, Array(Row(9, 9, 9)))
+      }
+    }
+  }
 }
 
 case class Foo(bar: Option[String])