[SPARK-25331][SS] Make FileStreamSink ignore partitions of batches that have already been written to file system

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -973,7 +973,7 @@ object SQLConf {
     buildConf("spark.sql.streaming.commitProtocolClass")
       .internal()
       .stringConf
-      .createWithDefault("org.apache.spark.sql.execution.streaming.ManifestFileCommitProtocol")
+      .createWithDefault("org.apache.spark.sql.execution.streaming.StagingFileCommitProtocol")
 
   val STREAMING_MULTIPLE_WATERMARK_POLICY =
     buildConf("spark.sql.streaming.multipleWatermarkPolicy")

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

@@ -218,7 +218,7 @@ object FileFormatWriter extends Logging {
       hadoopConf.set("mapreduce.task.id", taskAttemptId.getTaskID.toString)
       hadoopConf.set("mapreduce.task.attempt.id", taskAttemptId.toString)
       hadoopConf.setBoolean("mapreduce.task.ismap", true)
-      hadoopConf.setInt("mapreduce.task.partition", 0)
+      hadoopConf.setInt("mapreduce.task.partition", sparkPartitionId)
 
       new TaskAttemptContextImpl(hadoopConf, taskAttemptId)
     }

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/FileStreamSink.scala

@@ -113,7 +113,7 @@ class FileStreamSink(
         outputPath = path)
 
       committer match {
-        case manifestCommitter: ManifestFileCommitProtocol =>
+        case manifestCommitter: ManifestCommitProtocol =>
           manifestCommitter.setupManifestOptions(fileLog, batchId)
         case _ =>  // Do nothing
       }

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ManifestCommitProtocol.scala

@@ -0,0 +1,32 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.streaming
+
+trait ManifestCommitProtocol {
+  @transient protected var fileLog: MetadataLog[Array[SinkFileStatus]] = _
+  protected var batchId: Long = _
+  /**
+   * Sets up the manifest log output and the batch id for this job.
+   * Must be called before any other function.
+   */
+  def setupManifestOptions(fileLog: MetadataLog[Array[SinkFileStatus]], batchId: Long): Unit = {
+    this.fileLog = fileLog
+    this.batchId = batchId
+  }
+
+}

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ManifestFileCommitProtocol.scala

@@ -35,23 +35,12 @@ import org.apache.spark.internal.io.FileCommitProtocol.TaskCommitMessage
  * @param path path to write the final output to.
  */
 class ManifestFileCommitProtocol(jobId: String, path: String)
-  extends FileCommitProtocol with Serializable with Logging {
+  extends FileCommitProtocol with Serializable with Logging
+  with ManifestCommitProtocol {
 
   // Track the list of files added by a task, only used on the executors.
   @transient private var addedFiles: ArrayBuffer[String] = _
 
-  @transient private var fileLog: FileStreamSinkLog = _
-  private var batchId: Long = _
-
-  /**
-   * Sets up the manifest log output and the batch id for this job.
-   * Must be called before any other function.
-   */
-  def setupManifestOptions(fileLog: FileStreamSinkLog, batchId: Long): Unit = {
-    this.fileLog = fileLog
-    this.batchId = batchId
-  }
-
   override def setupJob(jobContext: JobContext): Unit = {
     require(fileLog != null, "setupManifestOptions must be called before this function")
     // Do nothing

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StagingFileCommitProtocol.scala

@@ -0,0 +1,141 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.streaming
+
+import org.apache.hadoop.fs.{FileAlreadyExistsException, FileContext, Path}
+import org.apache.hadoop.mapreduce.{JobContext, TaskAttemptContext}
+
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.io.FileCommitProtocol
+import org.apache.spark.internal.io.FileCommitProtocol.TaskCommitMessage
+
+class StagingFileCommitProtocol(jobId: String, path: String)
+  extends FileCommitProtocol with Serializable with Logging
+  with ManifestCommitProtocol {
+  private var stagingDir: Option[Path] = None
+
+
+  def jobStagingDir: Path = {
+    new Path(new Path(path, "staging"), s"job-$jobId")
+  }
+
+  override def setupJob(jobContext: JobContext): Unit = {
+    jobStagingDir.getFileSystem(jobContext.getConfiguration).delete(jobStagingDir, true)
+    logInfo(s"Job $jobId set up")
+  }
+
+
+  override def setupTask(taskContext: TaskAttemptContext): Unit = {
+    stagingDir = Some(new Path(jobStagingDir, s"partition-${partition(taskContext)}"))
+    stagingDir.get.getFileSystem(taskContext.getConfiguration).delete(stagingDir.get, true)
+    logInfo(s"Task set up to handle partition ${partition(taskContext)} in job $jobId")
+
+  }
+
+  private def partition(taskContext: TaskAttemptContext) = {
+    taskContext.getConfiguration.getInt("mapreduce.task.partition", -1)
+  }
+
+
+  override def commitJob(jobContext: JobContext, taskCommits: Seq[TaskCommitMessage]): Unit = {
+    val fs = jobStagingDir.getFileSystem(jobContext.getConfiguration)
+    val fileCtx = FileContext.getFileContext
+
+    def moveIfPossible(next: Path, target: Path) = {
+      try {
+        fileCtx.rename(next, target)
+      } catch {
+        case _: FileAlreadyExistsException =>
+          val status = fileCtx.getFileStatus(target)
+          logWarning(s"File ${target.toUri.toASCIIString} has already been generated " +
+            s"earlier (${status.toString}), deleting instead of moving " +
+            s"recently generated file: ${fileCtx.getFileStatus(target).toString}")
+          fileCtx.delete(next, false)
+      }
+    }
+
+    def moveEach(from: Path, to: String) = {
+      val files = fs.listFiles(from, true)
+      val statuses = Array.newBuilder[SinkFileStatus]
+      while (files.hasNext) {
+        val next = files.next().getPath
+        val target = if (next.getParent.getName.startsWith(outputPartitionPrefix)) {
+          val subdir = next.getParent.getName.substring(outputPartitionPrefix.length)
+              .replaceAll(subdirEscapeSequence, "/")
+          val outputPartition = new Path(to, subdir)
+          fs.mkdirs(outputPartition)
+          new Path(outputPartition, next.getName)
+        } else {
+          new Path(to, next.getName)
+        }
+        moveIfPossible(next, target)
+        statuses += SinkFileStatus(fs.getFileStatus(target))
+      }
+      if (fileLog.add(batchId, statuses.result)) {
+        logInfo(s"Job $jobId committed")
+      } else {
+        throw new IllegalStateException(s"Race while writing batch $batchId")
+      }
+    }
+
+    moveEach(jobStagingDir, path)
+
+    Seq()
+  }
+
+  override def abortJob(jobContext: JobContext): Unit = {}
+
+
+  private var fileCounter: Int = -1
+
+  private def nextCounter: Int = {
+    fileCounter += 1
+    fileCounter
+  }
+
+  private val outputPartitionPrefix = "part_prefix_"
+
+  private val subdirEscapeSequence = "___per___"
+
+  override def newTaskTempFile(
+      taskContext: TaskAttemptContext, dir: Option[String], ext: String): String = {
+    val staging = stagingDir.getOrElse(
+      throw new IllegalStateException("Staging dir needs to be initilized in setupTask()"))
+    val targetDir = dir.map(d => new Path(staging, stagingReplacementDir(d))).getOrElse(staging)
+    val res = new Path(targetDir, s"part-j$jobId-p${partition(taskContext)}-c$nextCounter$ext")
+      .toString
+    logInfo(s"New file generated $res")
+    res
+  }
+
+  private def stagingReplacementDir(d: String) = {
+    outputPartitionPrefix + d.replaceAll("/", subdirEscapeSequence)
+  }
+
+  override def newTaskTempFileAbsPath(
+      taskContext: TaskAttemptContext, absoluteDir: String, ext: String): String = {
+    throw new UnsupportedOperationException(
+      s"$this does not support adding files with an absolute path")
+  }
+
+  override def commitTask(taskContext: TaskAttemptContext): TaskCommitMessage = {
+    new TaskCommitMessage(None)
+  }
+
+  override def abortTask(taskContext: TaskAttemptContext): Unit = {}
+}