[SPARK-23202][SQL] Break down DataSourceV2Writer.commit into two phase

external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaStreamWriter.scala

@@ -49,8 +49,9 @@ class KafkaStreamWriter(
   override def createInternalRowWriterFactory(): KafkaStreamWriterFactory =
     KafkaStreamWriterFactory(topic, producerParams, schema)
 
-  override def commit(epochId: Long, messages: Array[WriterCommitMessage]): Unit = {}
-  override def abort(epochId: Long, messages: Array[WriterCommitMessage]): Unit = {}
+  override def add(message: WriterCommitMessage): Unit = {}
+  override def commit(epochId: Long): Unit = {}
+  override def abort(epochId: Long): Unit = {}
 }
 
 /**

sql/core/src/main/java/org/apache/spark/sql/sources/v2/streaming/writer/StreamWriter.java

@@ -32,40 +32,44 @@
 @InterfaceStability.Evolving
 public interface StreamWriter extends DataSourceWriter {
   /**
-   * Commits this writing job for the specified epoch with a list of commit messages. The commit
-   * messages are collected from successful data writers and are produced by
-   * {@link DataWriter#commit()}.
+   * Commits this writing job for the specified epoch.
    *
-   * If this method fails (by throwing an exception), this writing job is considered to have been
-   * failed, and the execution engine will attempt to call {@link #abort(WriterCommitMessage[])}.
+   * When this method is called, the number of commit messages added by
+   * {@link #add(WriterCommitMessage)} equals to the number of input data partitions.
+   *
+   * If this method fails (by throwing an exception), this writing job is considered to to have been
+   * failed, and {@link #abort()} would be called. The state of the destination
+   * is undefined and @{@link #abort()} may not be able to deal with it.
    *
    * To support exactly-once processing, writer implementations should ensure that this method is
    * idempotent. The execution engine may call commit() multiple times for the same epoch
    * in some circumstances.
    */
-  void commit(long epochId, WriterCommitMessage[] messages);
+  void commit(long epochId);
 
   /**
-   * Aborts this writing job because some data writers are failed and keep failing when retry, or
-   * the Spark job fails with some unknown reasons, or {@link #commit(WriterCommitMessage[])} fails.
+   * Aborts this writing job because some data writers are failed and keep failing when retry,
+   * or the Spark job fails with some unknown reasons,
+   * or {@link #commit()} / {@link #add(WriterCommitMessage)} fails
    *
    * If this method fails (by throwing an exception), the underlying data source may require manual
    * cleanup.
    *
-   * Unless the abort is triggered by the failure of commit, the given messages should have some
-   * null slots as there maybe only a few data writers that are committed before the abort
-   * happens, or some data writers were committed but their commit messages haven't reached the
-   * driver when the abort is triggered. So this is just a "best effort" for data sources to
-   * clean up the data left by data writers.
+   * Unless the abort is triggered by the failure of commit, the number of commit
+   * messages added by {@link #add(WriterCommitMessage)} should be smaller than the number
+   * of input data partitions, as there may be only a few data writers that are committed
+   * before the abort happens, or some data writers were committed but their commit messages
+   * haven't reached the driver when the abort is triggered. So this is just a "best effort"
+   * for data sources to clean up the data left by data writers.
    */
-  void abort(long epochId, WriterCommitMessage[] messages);
+  void abort(long epochId);
 
-  default void commit(WriterCommitMessage[] messages) {
+  default void commit() {
     throw new UnsupportedOperationException(
         "Commit without epoch should not be called with StreamWriter");
   }
 
-  default void abort(WriterCommitMessage[] messages) {
+  default void abort() {
     throw new UnsupportedOperationException(
         "Abort without epoch should not be called with StreamWriter");
   }

sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataSourceWriter.java

@@ -40,16 +40,21 @@
  *   1. Create a writer factory by {@link #createWriterFactory()}, serialize and send it to all the
  *      partitions of the input data(RDD).
  *   2. For each partition, create the data writer, and write the data of the partition with this
- *      writer. If all the data are written successfully, call {@link DataWriter#commit()}. If
- *      exception happens during the writing, call {@link DataWriter#abort()}.
- *   3. If all writers are successfully committed, call {@link #commit(WriterCommitMessage[])}. If
- *      some writers are aborted, or the job failed with an unknown reason, call
- *      {@link #abort(WriterCommitMessage[])}.
+ *      writer. If one data writer finishes successfully, the commit message will be sent back to
+ *      the driver side and Spark will call {@link #add(WriterCommitMessage)}.
+ *      If exception happens during the writing, call {@link DataWriter#abort()}.
+ *   3. If all the data writers finish successfully, and {@link #add(WriterCommitMessage)} is
+ *      successfully called for all the commit messages, Spark will call {@link #commit()}.
+ *      If any of the data writers failed, or any of the {@link #add(WriterCommitMessage)}
+ *      calls failed, or the job failed with an unknown reason, call {@link #abort()}.
  *
  * While Spark will retry failed writing tasks, Spark won't retry failed writing jobs. Users should
  * do it manually in their Spark applications if they want to retry.
  *
- * Please refer to the documentation of commit/abort methods for detailed specifications.
+ * All these methods are guaranteed to be called in a single thread on driver side.
+ * No concurrency control is needed.
+ *
+ * Please refer to the documentation of add/commit/abort methods for detailed specifications.
  */
 @InterfaceStability.Evolving
 public interface DataSourceWriter {
@@ -63,32 +68,42 @@ public interface DataSourceWriter {
   DataWriterFactory<Row> createWriterFactory();
 
   /**
-   * Commits this writing job with a list of commit messages. The commit messages are collected from
-   * successful data writers and are produced by {@link DataWriter#commit()}.
+   * Handles a commit message which is collected from a successful data writer.
+   *
+   * Note that, implementations might need to cache all commit messages before calling
+   * {@link #commit()} or {@link #abort()}.
    *
    * If this method fails (by throwing an exception), this writing job is considered to to have been
-   * failed, and {@link #abort(WriterCommitMessage[])} would be called. The state of the destination
-   * is undefined and @{@link #abort(WriterCommitMessage[])} may not be able to deal with it.
+   * failed, and {@link #abort()} would be called. The state of the destination
+   * is undefined and @{@link #abort()} may not be able to deal with it.
+   */
+  void add(WriterCommitMessage message);
+
+  /**
+   * Commits this writing job.
+   * When this method is called, the number of commit messages added by
+   * {@link #add(WriterCommitMessage)} equals to the number of input data partitions.
    *
-   * Note that, one partition may have multiple committed data writers because of speculative tasks.
-   * Spark will pick the first successful one and get its commit message. Implementations should be
-   * aware of this and handle it correctly, e.g., have a coordinator to make sure only one data
-   * writer can commit, or have a way to clean up the data of already-committed writers.
+   * If this method fails (by throwing an exception), this writing job is considered to to have been
+   * failed, and {@link #abort()} would be called. The state of the destination
+   * is undefined and @{@link #abort()} may not be able to deal with it.
    */
-  void commit(WriterCommitMessage[] messages);
+  void commit();
 
   /**
-   * Aborts this writing job because some data writers are failed and keep failing when retry, or
-   * the Spark job fails with some unknown reasons, or {@link #commit(WriterCommitMessage[])} fails.
+   * Aborts this writing job because some data writers are failed and keep failing when retry,
+   * or the Spark job fails with some unknown reasons,
+   * or {@link #commit()} / {@link #add(WriterCommitMessage)} fails.
    *
    * If this method fails (by throwing an exception), the underlying data source may require manual
    * cleanup.
    *
-   * Unless the abort is triggered by the failure of commit, the given messages should have some
-   * null slots as there maybe only a few data writers that are committed before the abort
-   * happens, or some data writers were committed but their commit messages haven't reached the
-   * driver when the abort is triggered. So this is just a "best effort" for data sources to
-   * clean up the data left by data writers.
+   * Unless the abort is triggered by the failure of {@link #commit()}, the number of commit
+   * messages added by {@link #add(WriterCommitMessage)} should be smaller than the number
+   * of input data partitions, as there may be only a few data writers that are committed
+   * before the abort happens, or some data writers were committed but their commit messages
+   * haven't reached the driver when the abort is triggered. So this is just a "best effort"
+   * for data sources to clean up the data left by data writers.
    */
-  void abort(WriterCommitMessage[] messages);
+  void abort();
 }

sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataWriter.java

@@ -33,11 +33,11 @@
  *
  * If this data writer succeeds(all records are successfully written and {@link #commit()}
  * succeeds), a {@link WriterCommitMessage} will be sent to the driver side and pass to
- * {@link DataSourceWriter#commit(WriterCommitMessage[])} with commit messages from other data
+ * {@link DataSourceWriter#commit()} with commit messages from other data
  * writers. If this data writer fails(one record fails to write or {@link #commit()} fails), an
  * exception will be sent to the driver side, and Spark will retry this writing task for some times,
  * each time {@link DataWriterFactory#createDataWriter(int, int)} gets a different `attemptNumber`,
- * and finally call {@link DataSourceWriter#abort(WriterCommitMessage[])} if all retry fail.
+ * and finally call {@link DataSourceWriter#abort()} if all retry fail.
  *
  * Besides the retry mechanism, Spark may launch speculative tasks if the existing writing task
  * takes too long to finish. Different from retried tasks, which are launched one by one after the
@@ -69,10 +69,10 @@ public interface DataWriter<T> {
   /**
    * Commits this writer after all records are written successfully, returns a commit message which
    * will be sent back to driver side and passed to
-   * {@link DataSourceWriter#commit(WriterCommitMessage[])}.
+   * {@link DataSourceWriter#commit()}.
    *
    * The written data should only be visible to data source readers after
-   * {@link DataSourceWriter#commit(WriterCommitMessage[])} succeeds, which means this method
+   * {@link DataSourceWriter#commit()} succeeds, which means this method
    * should still "hide" the written data and ask the {@link DataSourceWriter} at driver side to
    * do the final commit via {@link WriterCommitMessage}.
    *
@@ -91,7 +91,7 @@ public interface DataWriter<T> {
    * failed.
    *
    * If this method fails(by throwing an exception), the underlying data source may have garbage
-   * that need to be cleaned by {@link DataSourceWriter#abort(WriterCommitMessage[])} or manually,
+   * that need to be cleaned by {@link DataSourceWriter#abort()} or manually,
    * but these garbage should not be visible to data source readers.
    *
    * @throws IOException if failure happens during disk/network IO like writing files.

sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/WriterCommitMessage.java

@@ -23,10 +23,10 @@
 
 /**
  * A commit message returned by {@link DataWriter#commit()} and will be sent back to the driver side
- * as the input parameter of {@link DataSourceWriter#commit(WriterCommitMessage[])}.
+ * as the input parameter of {@link DataSourceWriter#commit()}.
  *
  * This is an empty interface, data sources should define their own message class and use it in
- * their {@link DataWriter#commit()} and {@link DataSourceWriter#commit(WriterCommitMessage[])}
+ * their {@link DataWriter#commit()} and {@link DataSourceWriter#commit()}
  * implementations.
  */
 @InterfaceStability.Evolving

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

@@ -54,10 +54,9 @@ case class WriteToDataSourceV2Exec(writer: DataSourceWriter, query: SparkPlan) e
     }
 
     val rdd = query.execute()
-    val messages = new Array[WriterCommitMessage](rdd.partitions.length)
 
     logInfo(s"Start processing data source writer: $writer. " +
-      s"The input RDD has ${messages.length} partitions.")
+      s"The input RDD has ${rdd.partitions.length} partitions.")
 
     try {
       val runTask = writer match {
@@ -80,12 +79,12 @@ case class WriteToDataSourceV2Exec(writer: DataSourceWriter, query: SparkPlan) e
         rdd,
         runTask,
         rdd.partitions.indices,
-        (index, message: WriterCommitMessage) => messages(index) = message
+        (_, message: WriterCommitMessage) => writer.add(message)
       )
 
       if (!writer.isInstanceOf[StreamWriter]) {
         logInfo(s"Data source writer $writer is committing.")
-        writer.commit(messages)
+        writer.commit()
         logInfo(s"Data source writer $writer committed.")
       }
     } catch {
@@ -94,7 +93,7 @@ case class WriteToDataSourceV2Exec(writer: DataSourceWriter, query: SparkPlan) e
       case cause: Throwable =>
         logError(s"Data source writer $writer is aborting.")
         try {
-          writer.abort(messages)
+          writer.abort()
         } catch {
           case t: Throwable =>
             logError(s"Data source writer $writer failed to abort.")

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

@@ -148,7 +148,7 @@ private[continuous] class EpochCoordinator(
       logDebug(s"Epoch $epoch has received commits from all partitions. Committing globally.")
       // Sequencing is important here. We must commit to the writer before recording the commit
       // in the query, or we will end up dropping the commit if we restart in the middle.
-      writer.commit(epoch, thisEpochCommits.toArray)
+      writer.commit(epoch)
       query.commit(epoch)
 
       // Cleanup state from before this epoch, now that we know all partitions are forever past it.
@@ -170,6 +170,7 @@ private[continuous] class EpochCoordinator(
       logDebug(s"Got commit from partition $partitionId at epoch $epoch: $message")
       if (!partitionCommits.isDefinedAt((epoch, partitionId))) {
         partitionCommits.put((epoch, partitionId), message)
+        writer.add(message)
         resolveCommitsAtEpoch(epoch)
       }
 

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

@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.execution.streaming.sources
 
+import scala.collection.mutable.ArrayBuffer
+
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.{Row, SparkSession}
 import org.apache.spark.sql.sources.v2.DataSourceOptions
@@ -39,13 +41,22 @@ class ConsoleWriter(schema: StructType, options: DataSourceOptions)
 
   def createWriterFactory(): DataWriterFactory[Row] = PackedRowWriterFactory
 
-  override def commit(epochId: Long, messages: Array[WriterCommitMessage]): Unit = {
+  private val messages = new ArrayBuffer[WriterCommitMessage]()
+
+  override def add(message: WriterCommitMessage): Unit = {
+    messages += message
+  }
+
+  override def commit(epochId: Long): Unit = {
     // We have to print a "Batch" label for the epoch for compatibility with the pre-data source V2
     // behavior.
-    printRows(messages, schema, s"Batch: $epochId")
+    printRows(messages.toArray, schema, s"Batch: $epochId")
+    messages.clear()
   }
 
-  def abort(epochId: Long, messages: Array[WriterCommitMessage]): Unit = {}
+  def abort(epochId: Long): Unit = {
+    messages.clear()
+  }
 
   protected def printRows(
       commitMessages: Array[WriterCommitMessage],

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

@@ -28,22 +28,30 @@ import org.apache.spark.sql.sources.v2.writer.{DataSourceWriter, DataWriterFacto
  * streaming writer.
  */
 class MicroBatchWriter(batchId: Long, writer: StreamWriter) extends DataSourceWriter {
-  override def commit(messages: Array[WriterCommitMessage]): Unit = {
-    writer.commit(batchId, messages)
+  override def add(message: WriterCommitMessage): Unit = {
+    writer.add(message)
   }
 
-  override def abort(messages: Array[WriterCommitMessage]): Unit = writer.abort(batchId, messages)
+  override def commit(): Unit = {
+    writer.commit(batchId)
+  }
+
+  override def abort(): Unit = writer.abort(batchId)
 
   override def createWriterFactory(): DataWriterFactory[Row] = writer.createWriterFactory()
 }
 
 class InternalRowMicroBatchWriter(batchId: Long, writer: StreamWriter)
   extends DataSourceWriter with SupportsWriteInternalRow {
-  override def commit(messages: Array[WriterCommitMessage]): Unit = {
-    writer.commit(batchId, messages)
+  override def add(message: WriterCommitMessage): Unit = {
+    writer.add(message)
+  }
+
+  override def commit(): Unit = {
+    writer.commit(batchId)
   }
 
-  override def abort(messages: Array[WriterCommitMessage]): Unit = writer.abort(batchId, messages)
+  override def abort(): Unit = writer.abort(batchId)
 
   override def createInternalRowWriterFactory(): DataWriterFactory[InternalRow] =
     writer match {

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

@@ -118,15 +118,23 @@ class MemoryWriter(sink: MemorySinkV2, batchId: Long, outputMode: OutputMode)
 
   override def createWriterFactory: MemoryWriterFactory = MemoryWriterFactory(outputMode)
 
-  def commit(messages: Array[WriterCommitMessage]): Unit = {
+  private val messages = new ArrayBuffer[WriterCommitMessage]()
+
+  override def add(message: WriterCommitMessage): Unit = {
+    messages += message
+  }
+
+  def commit(): Unit = {
     val newRows = messages.flatMap {
       case message: MemoryWriterCommitMessage => message.data
-    }
+    }.toArray
     sink.write(batchId, outputMode, newRows)
+    messages.clear()
   }
 
-  override def abort(messages: Array[WriterCommitMessage]): Unit = {
+  override def abort(): Unit = {
     // Don't accept any of the new input.
+    messages.clear()
   }
 }
 
@@ -135,15 +143,23 @@ class MemoryStreamWriter(val sink: MemorySinkV2, outputMode: OutputMode)
 
   override def createWriterFactory: MemoryWriterFactory = MemoryWriterFactory(outputMode)
 
-  override def commit(epochId: Long, messages: Array[WriterCommitMessage]): Unit = {
+  private val messages = new ArrayBuffer[WriterCommitMessage]()
+
+  override def add(message: WriterCommitMessage): Unit = {
+    messages += message
+  }
+
+  override def commit(epochId: Long): Unit = {
     val newRows = messages.flatMap {
       case message: MemoryWriterCommitMessage => message.data
-    }
+    }.toArray
     sink.write(epochId, outputMode, newRows)
+    messages.clear()
   }
 
-  override def abort(epochId: Long, messages: Array[WriterCommitMessage]): Unit = {
+  override def abort(epochId: Long): Unit = {
     // Don't accept any of the new input.
+    messages.clear()
   }
 }