From 7ae1029bb77636a96e59002a08e6a649ddea0611 Mon Sep 17 00:00:00 2001
From: Wang Gengliang <ltnwgl@gmail.com>
Date: Wed, 24 Jan 2018 17:15:38 +0800
Subject: [PATCH 01/10] add api 'add' in DataSourceV2Writer

---
 .../v2/streaming/writer/StreamWriter.java     | 16 +++---
 .../sources/v2/writer/DataSourceWriter.java   | 44 ++++++++--------
 .../sql/sources/v2/writer/DataWriter.java     | 10 ++--
 .../v2/writer/WriterCommitMessage.java        |  4 +-
 .../datasources/v2/WriteToDataSourceV2.scala  | 10 ++--
 .../continuous/EpochCoordinator.scala         |  3 +-
 .../streaming/sources/memoryV2.scala          | 26 +++++++---
 .../streaming/MemorySinkV2Suite.scala         | 52 +++++++++++--------
 .../sources/v2/SimpleWritableDataSource.scala |  6 ++-
 9 files changed, 96 insertions(+), 75 deletions(-)

diff --git a/sql/core/src/main/java/org/apache/spark/sql/sources/v2/streaming/writer/StreamWriter.java b/sql/core/src/main/java/org/apache/spark/sql/sources/v2/streaming/writer/StreamWriter.java
index 915ee6c4fb390..854b0a2d158e5 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/sources/v2/streaming/writer/StreamWriter.java
+++ b/sql/core/src/main/java/org/apache/spark/sql/sources/v2/streaming/writer/StreamWriter.java
@@ -37,18 +37,18 @@ public interface StreamWriter extends DataSourceWriter {
    * {@link DataWriter#commit()}.
    *
    * 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[])}.
+   * failed, and the execution engine will attempt to call {@link #abort()}.
    *
    * 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.
    *
@@ -58,14 +58,14 @@ public interface StreamWriter extends DataSourceWriter {
    * 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");
   }
diff --git a/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataSourceWriter.java b/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataSourceWriter.java
index d89d27d0e5b1b..be7d741402c5f 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataSourceWriter.java
+++ b/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataSourceWriter.java
@@ -40,11 +40,13 @@
  *   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
+ *      writer. If all the data are written successfully, call {@link DataWriter#commit()}.
+ *      On a writer being successfully committed, call {@link #add(WriterCommitMessage)} to
+ *      handle its commit message.
+ *      If exception happens during the writing, call {@link DataWriter#abort()}.
+ *   3. If all writers are successfully committed, call {@link #commit()}. If
  *      some writers are aborted, or the job failed with an unknown reason, call
- *      {@link #abort(WriterCommitMessage[])}.
+ *      {@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.
@@ -63,32 +65,30 @@ 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 produced by {@link DataWriter#commit()}.
    *
    * 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.
    *
-   * 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.
    */
-  void abort(WriterCommitMessage[] messages);
+  void abort();
 }
diff --git a/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataWriter.java b/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataWriter.java
index 53941a89ba94e..26dc8310be3df 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataWriter.java
+++ b/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.
diff --git a/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/WriterCommitMessage.java b/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/WriterCommitMessage.java
index 9e38836c0edf9..c082054be6966 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/WriterCommitMessage.java
+++ b/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
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2.scala
index c544adbf32cdf..12e9cf3bab659 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2.scala
@@ -54,10 +54,6 @@ 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.")
 
     try {
       val runTask = writer match {
@@ -80,12 +76,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 +90,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.")
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/EpochCoordinator.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/EpochCoordinator.scala
index 84d262116cb46..891a10bdb1509 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/EpochCoordinator.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/EpochCoordinator.scala
@@ -148,7 +148,8 @@ 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)
+      thisEpochCommits.foreach(writer.add(_))
+      writer.commit(epoch)
       query.commit(epoch)
 
       // Cleanup state from before this epoch, now that we know all partitions are forever past it.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/sources/memoryV2.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/sources/memoryV2.scala
index 58767261dc684..1b6d1f0b89145 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/sources/memoryV2.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/sources/memoryV2.scala
@@ -118,14 +118,21 @@ 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 = synchronized {
+    messages += message
+  }
+
+  def commit(): Unit = synchronized {
     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.
   }
 }
@@ -135,14 +142,21 @@ 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 = synchronized {
+    messages += message
+  }
+
+  override def commit(epochId: Long): Unit = synchronized {
     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.
   }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/MemorySinkV2Suite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/MemorySinkV2Suite.scala
index 9be22d94b5654..5dd2b447a9e36 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/MemorySinkV2Suite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/MemorySinkV2Suite.scala
@@ -41,19 +41,22 @@ class MemorySinkV2Suite extends StreamTest with BeforeAndAfter {
   test("continuous writer") {
     val sink = new MemorySinkV2
     val writer = new MemoryStreamWriter(sink, OutputMode.Append())
-    writer.commit(0,
-      Array(
-        MemoryWriterCommitMessage(0, Seq(Row(1), Row(2))),
-        MemoryWriterCommitMessage(1, Seq(Row(3), Row(4))),
-        MemoryWriterCommitMessage(2, Seq(Row(6), Row(7)))
-      ))
+    val messages = Seq(
+      MemoryWriterCommitMessage(0, Seq(Row(1), Row(2))),
+      MemoryWriterCommitMessage(1, Seq(Row(3), Row(4))),
+      MemoryWriterCommitMessage(2, Seq(Row(6), Row(7)))
+    )
+    messages.foreach(writer.add(_))
+    writer.commit(0)
     assert(sink.latestBatchId.contains(0))
     assert(sink.latestBatchData.map(_.getInt(0)).sorted == Seq(1, 2, 3, 4, 6, 7))
-    writer.commit(19,
-      Array(
-        MemoryWriterCommitMessage(3, Seq(Row(11), Row(22))),
-        MemoryWriterCommitMessage(0, Seq(Row(33)))
-      ))
+
+    val newMessages = Seq(
+      MemoryWriterCommitMessage(3, Seq(Row(11), Row(22))),
+      MemoryWriterCommitMessage(0, Seq(Row(33)))
+    )
+    newMessages.foreach(writer.add(_))
+    writer.commit(19)
     assert(sink.latestBatchId.contains(19))
     assert(sink.latestBatchData.map(_.getInt(0)).sorted == Seq(11, 22, 33))
 
@@ -62,19 +65,24 @@ class MemorySinkV2Suite extends StreamTest with BeforeAndAfter {
 
   test("microbatch writer") {
     val sink = new MemorySinkV2
-    new MemoryWriter(sink, 0, OutputMode.Append()).commit(
-      Array(
-        MemoryWriterCommitMessage(0, Seq(Row(1), Row(2))),
-        MemoryWriterCommitMessage(1, Seq(Row(3), Row(4))),
-        MemoryWriterCommitMessage(2, Seq(Row(6), Row(7)))
-      ))
+    val writer = new MemoryWriter(sink, 0, OutputMode.Append())
+    val messages = Seq(
+      MemoryWriterCommitMessage(0, Seq(Row(1), Row(2))),
+      MemoryWriterCommitMessage(1, Seq(Row(3), Row(4))),
+      MemoryWriterCommitMessage(2, Seq(Row(6), Row(7)))
+    )
+    messages.foreach(writer.add(_))
+    writer.commit()
     assert(sink.latestBatchId.contains(0))
     assert(sink.latestBatchData.map(_.getInt(0)).sorted == Seq(1, 2, 3, 4, 6, 7))
-    new MemoryWriter(sink, 19, OutputMode.Append()).commit(
-      Array(
-        MemoryWriterCommitMessage(3, Seq(Row(11), Row(22))),
-        MemoryWriterCommitMessage(0, Seq(Row(33)))
-      ))
+
+    val newWriter = new MemoryWriter(sink, 19, OutputMode.Append())
+    val newMessages = Seq(
+      MemoryWriterCommitMessage(3, Seq(Row(11), Row(22))),
+      MemoryWriterCommitMessage(0, Seq(Row(33)))
+    )
+    newMessages.foreach(newWriter.add(_))
+    newWriter.commit()
     assert(sink.latestBatchId.contains(19))
     assert(sink.latestBatchData.map(_.getInt(0)).sorted == Seq(11, 22, 33))
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/v2/SimpleWritableDataSource.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/v2/SimpleWritableDataSource.scala
index a131b16953e3b..5c5a1429088b3 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/sources/v2/SimpleWritableDataSource.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/v2/SimpleWritableDataSource.scala
@@ -69,7 +69,9 @@ class SimpleWritableDataSource extends DataSourceV2 with ReadSupport with WriteS
       new SimpleCSVDataWriterFactory(path, jobId, new SerializableConfiguration(conf))
     }
 
-    override def commit(messages: Array[WriterCommitMessage]): Unit = {
+    override def add(message: WriterCommitMessage): Unit = {}
+
+    override def commit(): Unit = {
       val finalPath = new Path(path)
       val jobPath = new Path(new Path(finalPath, "_temporary"), jobId)
       val fs = jobPath.getFileSystem(conf)
@@ -85,7 +87,7 @@ class SimpleWritableDataSource extends DataSourceV2 with ReadSupport with WriteS
       }
     }
 
-    override def abort(messages: Array[WriterCommitMessage]): Unit = {
+    override def abort(): Unit = {
       val jobPath = new Path(new Path(path, "_temporary"), jobId)
       val fs = jobPath.getFileSystem(conf)
       fs.delete(jobPath, true)

From 0d2e39a0413791307ff2c86678ac6d5435dc42cd Mon Sep 17 00:00:00 2001
From: Wang Gengliang <ltnwgl@gmail.com>
Date: Thu, 25 Jan 2018 20:04:52 +0800
Subject: [PATCH 02/10] Fix with latest code

---
 .../spark/sql/kafka010/KafkaStreamWriter.scala    |  5 +++--
 .../streaming/sources/ConsoleWriter.scala         | 15 ++++++++++++---
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaStreamWriter.scala b/external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaStreamWriter.scala
index a24efdefa4464..d809ed6b952b3 100644
--- a/external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaStreamWriter.scala
+++ b/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 = {}
 }
 
 /**
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/sources/ConsoleWriter.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/sources/ConsoleWriter.scala
index d46f4d7b86360..96d8e89f2570e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/sources/ConsoleWriter.scala
+++ b/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,20 @@ 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 = synchronized {
+    messages += message
+  }
+
+  override def commit(epochId: Long): Unit = synchronized {
     // 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 = {}
 
   protected def printRows(
       commitMessages: Array[WriterCommitMessage],

From 59b1857fa105e6e65ab40940dcc172b37d7f7fd0 Mon Sep 17 00:00:00 2001
From: Wang Gengliang <ltnwgl@gmail.com>
Date: Fri, 26 Jan 2018 14:56:26 +0800
Subject: [PATCH 03/10] Fix ConsoleWriterSuite

---
 .../sql/execution/streaming/sources/ConsoleWriterSuite.scala | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/sources/ConsoleWriterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/sources/ConsoleWriterSuite.scala
index 55acf2ba28d2f..e2be2ed6eaad8 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/sources/ConsoleWriterSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/sources/ConsoleWriterSuite.scala
@@ -45,7 +45,8 @@ class ConsoleWriterSuite extends StreamTest {
       }
     }
 
-    assert(captured.toString() ==
+    // The order of data in one batch can be random
+    assert(captured.toString().length ==
       """-------------------------------------------
         |Batch: 0
         |-------------------------------------------
@@ -76,7 +77,7 @@ class ConsoleWriterSuite extends StreamTest {
         |+-----+
         |+-----+
         |
-        |""".stripMargin)
+        |""".stripMargin.length)
   }
 
   test("microbatch - with numRows") {

From 5938ab37f2ad2c1a1665219027a77b1853893502 Mon Sep 17 00:00:00 2001
From: Wang Gengliang <ltnwgl@gmail.com>
Date: Fri, 26 Jan 2018 15:15:40 +0800
Subject: [PATCH 04/10] Better fix

---
 .../streaming/sources/ConsoleWriterSuite.scala       | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/sources/ConsoleWriterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/sources/ConsoleWriterSuite.scala
index e2be2ed6eaad8..2cbdbf34b44f2 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/sources/ConsoleWriterSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/sources/ConsoleWriterSuite.scala
@@ -19,14 +19,17 @@ package org.apache.spark.sql.execution.streaming.sources
 
 import java.io.ByteArrayOutputStream
 
-import org.scalatest.time.SpanSugar._
-
+import org.apache.spark.SparkConf
 import org.apache.spark.sql.execution.streaming.MemoryStream
 import org.apache.spark.sql.streaming.{StreamTest, Trigger}
 
 class ConsoleWriterSuite extends StreamTest {
   import testImplicits._
 
+  override def sparkConf: SparkConf = {
+    super.sparkConf.set("spark.default.parallelism", "1")
+  }
+
   test("microbatch - default") {
     val input = MemoryStream[Int]
 
@@ -45,8 +48,7 @@ class ConsoleWriterSuite extends StreamTest {
       }
     }
 
-    // The order of data in one batch can be random
-    assert(captured.toString().length ==
+    assert(captured.toString() ==
       """-------------------------------------------
         |Batch: 0
         |-------------------------------------------
@@ -77,7 +79,7 @@ class ConsoleWriterSuite extends StreamTest {
         |+-----+
         |+-----+
         |
-        |""".stripMargin.length)
+        |""".stripMargin)
   }
 
   test("microbatch - with numRows") {

From 939ad0608cb1a838b7256f972a44399ea463f884 Mon Sep 17 00:00:00 2001
From: Wang Gengliang <ltnwgl@gmail.com>
Date: Fri, 26 Jan 2018 17:26:22 +0800
Subject: [PATCH 05/10] Fix

---
 .../sources/ConsoleWriterSuite.scala          | 22 ++++++++-----------
 1 file changed, 9 insertions(+), 13 deletions(-)

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/sources/ConsoleWriterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/sources/ConsoleWriterSuite.scala
index 2cbdbf34b44f2..441bcc793dcdf 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/sources/ConsoleWriterSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/sources/ConsoleWriterSuite.scala
@@ -26,10 +26,6 @@ import org.apache.spark.sql.streaming.{StreamTest, Trigger}
 class ConsoleWriterSuite extends StreamTest {
   import testImplicits._
 
-  override def sparkConf: SparkConf = {
-    super.sparkConf.set("spark.default.parallelism", "1")
-  }
-
   test("microbatch - default") {
     val input = MemoryStream[Int]
 
@@ -37,9 +33,9 @@ class ConsoleWriterSuite extends StreamTest {
     Console.withOut(captured) {
       val query = input.toDF().writeStream.format("console").start()
       try {
-        input.addData(1, 2, 3)
+        input.addData(1, 1, 1)
         query.processAllAvailable()
-        input.addData(4, 5, 6)
+        input.addData(2, 2, 2)
         query.processAllAvailable()
         input.addData()
         query.processAllAvailable()
@@ -56,8 +52,8 @@ class ConsoleWriterSuite extends StreamTest {
         ||value|
         |+-----+
         ||    1|
-        ||    2|
-        ||    3|
+        ||    1|
+        ||    1|
         |+-----+
         |
         |-------------------------------------------
@@ -66,9 +62,9 @@ class ConsoleWriterSuite extends StreamTest {
         |+-----+
         ||value|
         |+-----+
-        ||    4|
-        ||    5|
-        ||    6|
+        ||    2|
+        ||    2|
+        ||    2|
         |+-----+
         |
         |-------------------------------------------
@@ -89,7 +85,7 @@ class ConsoleWriterSuite extends StreamTest {
     Console.withOut(captured) {
       val query = input.toDF().writeStream.format("console").option("NUMROWS", 2).start()
       try {
-        input.addData(1, 2, 3)
+        input.addData(1, 1, 1)
         query.processAllAvailable()
       } finally {
         query.stop()
@@ -104,7 +100,7 @@ class ConsoleWriterSuite extends StreamTest {
         ||value|
         |+-----+
         ||    1|
-        ||    2|
+        ||    1|
         |+-----+
         |only showing top 2 rows
         |

From aefef153db6f4d1e1d94696e6929a28fc03a8876 Mon Sep 17 00:00:00 2001
From: Wang Gengliang <ltnwgl@gmail.com>
Date: Tue, 30 Jan 2018 09:55:49 +0800
Subject: [PATCH 06/10] Fix MicroBatchWriter

---
 .../streaming/sources/MicroBatchWriter.scala  | 20 +++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/sources/MicroBatchWriter.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/sources/MicroBatchWriter.scala
index d7ce9a7b84479..f2ca9bd6fd344 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/sources/MicroBatchWriter.scala
+++ b/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 {

From 0e5ad8422a1e214eb5509be372b234f7e2604f72 Mon Sep 17 00:00:00 2001
From: Wang Gengliang <ltnwgl@gmail.com>
Date: Tue, 30 Jan 2018 15:28:22 +0800
Subject: [PATCH 07/10] Address comments

---
 .../sql/sources/v2/writer/DataSourceWriter.java | 17 ++++++++++-------
 .../streaming/continuous/EpochCoordinator.scala |  2 +-
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataSourceWriter.java b/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataSourceWriter.java
index be7d741402c5f..bd3a721cd9f0d 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataSourceWriter.java
+++ b/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataSourceWriter.java
@@ -40,13 +40,13 @@
  *   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()}.
- *      On a writer being successfully committed, call {@link #add(WriterCommitMessage)} to
- *      handle its commit message.
+ *      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 writers are successfully committed, call {@link #commit()}. If
- *      some writers are aborted, or the job failed with an unknown reason, call
- *      {@link #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.
@@ -65,7 +65,10 @@ public interface DataSourceWriter {
   DataWriterFactory<Row> createWriterFactory();
 
   /**
-   * Handles a commit message produced by {@link DataWriter#commit()}.
+   * Handles a commit message which is collected from a successful data writer in the executor side.
+   *
+   * 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()} would be called. The state of the destination
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/EpochCoordinator.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/EpochCoordinator.scala
index 891a10bdb1509..9b84c69e02fee 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/EpochCoordinator.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/EpochCoordinator.scala
@@ -148,7 +148,6 @@ 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.
-      thisEpochCommits.foreach(writer.add(_))
       writer.commit(epoch)
       query.commit(epoch)
 
@@ -171,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)
       }
 

From f72c86ce97ef7004c0a16b6fbe390308feda7759 Mon Sep 17 00:00:00 2001
From: Wang Gengliang <ltnwgl@gmail.com>
Date: Tue, 30 Jan 2018 20:25:13 +0800
Subject: [PATCH 08/10] address comments

---
 .../sources/v2/writer/DataSourceWriter.java   |  5 ++-
 .../datasources/v2/WriteToDataSourceV2.scala  |  3 ++
 .../streaming/sources/ConsoleWriter.scala     |  8 +++--
 .../streaming/sources/memoryV2.scala          | 10 +++---
 .../streaming/MemorySinkV2Suite.scala         | 32 ++++++-------------
 5 files changed, 28 insertions(+), 30 deletions(-)

diff --git a/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataSourceWriter.java b/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataSourceWriter.java
index bd3a721cd9f0d..4096f2e0b9784 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataSourceWriter.java
+++ b/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataSourceWriter.java
@@ -51,7 +51,10 @@
  * 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.
+ * In general, all these function calls should be thread-safe in driver side and there is no need
+ * to implement concurrency control.
+ *
+ * Please refer to the documentation of add/commit/abort methods for detailed specifications.
  */
 @InterfaceStability.Evolving
 public interface DataSourceWriter {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2.scala
index 12e9cf3bab659..6fe7e0a8c4338 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2.scala
@@ -55,6 +55,9 @@ case class WriteToDataSourceV2Exec(writer: DataSourceWriter, query: SparkPlan) e
 
     val rdd = query.execute()
 
+    logInfo(s"Start processing data source writer: $writer. " +
+      s"The input RDD has ${rdd.partitions.length} partitions.")
+
     try {
       val runTask = writer match {
         // This case means that we're doing continuous processing. In microbatch streaming, the
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/sources/ConsoleWriter.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/sources/ConsoleWriter.scala
index 96d8e89f2570e..308ace29d9b2d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/sources/ConsoleWriter.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/sources/ConsoleWriter.scala
@@ -43,18 +43,20 @@ class ConsoleWriter(schema: StructType, options: DataSourceOptions)
 
   private val messages = new ArrayBuffer[WriterCommitMessage]()
 
-  override def add(message: WriterCommitMessage): Unit = synchronized {
+  override def add(message: WriterCommitMessage): Unit = {
     messages += message
   }
 
-  override def commit(epochId: Long): Unit = synchronized {
+  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.toArray, schema, s"Batch: $epochId")
     messages.clear()
   }
 
-  def abort(epochId: Long): Unit = {}
+  def abort(epochId: Long): Unit = {
+    messages.clear()
+  }
 
   protected def printRows(
       commitMessages: Array[WriterCommitMessage],
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/sources/memoryV2.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/sources/memoryV2.scala
index 1b6d1f0b89145..3c4aeccb0f6b1 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/sources/memoryV2.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/sources/memoryV2.scala
@@ -120,11 +120,11 @@ class MemoryWriter(sink: MemorySinkV2, batchId: Long, outputMode: OutputMode)
 
   private val messages = new ArrayBuffer[WriterCommitMessage]()
 
-  override def add(message: WriterCommitMessage): Unit = synchronized {
+  override def add(message: WriterCommitMessage): Unit = {
     messages += message
   }
 
-  def commit(): Unit = synchronized {
+  def commit(): Unit = {
     val newRows = messages.flatMap {
       case message: MemoryWriterCommitMessage => message.data
     }.toArray
@@ -134,6 +134,7 @@ class MemoryWriter(sink: MemorySinkV2, batchId: Long, outputMode: OutputMode)
 
   override def abort(): Unit = {
     // Don't accept any of the new input.
+    messages.clear()
   }
 }
 
@@ -144,11 +145,11 @@ class MemoryStreamWriter(val sink: MemorySinkV2, outputMode: OutputMode)
 
   private val messages = new ArrayBuffer[WriterCommitMessage]()
 
-  override def add(message: WriterCommitMessage): Unit = synchronized {
+  override def add(message: WriterCommitMessage): Unit = {
     messages += message
   }
 
-  override def commit(epochId: Long): Unit = synchronized {
+  override def commit(epochId: Long): Unit = {
     val newRows = messages.flatMap {
       case message: MemoryWriterCommitMessage => message.data
     }.toArray
@@ -158,6 +159,7 @@ class MemoryStreamWriter(val sink: MemorySinkV2, outputMode: OutputMode)
 
   override def abort(epochId: Long): Unit = {
     // Don't accept any of the new input.
+    messages.clear()
   }
 }
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/MemorySinkV2Suite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/MemorySinkV2Suite.scala
index 5dd2b447a9e36..2d4e1d65478bb 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/MemorySinkV2Suite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/MemorySinkV2Suite.scala
@@ -41,21 +41,15 @@ class MemorySinkV2Suite extends StreamTest with BeforeAndAfter {
   test("continuous writer") {
     val sink = new MemorySinkV2
     val writer = new MemoryStreamWriter(sink, OutputMode.Append())
-    val messages = Seq(
-      MemoryWriterCommitMessage(0, Seq(Row(1), Row(2))),
-      MemoryWriterCommitMessage(1, Seq(Row(3), Row(4))),
-      MemoryWriterCommitMessage(2, Seq(Row(6), Row(7)))
-    )
-    messages.foreach(writer.add(_))
+    writer.add(MemoryWriterCommitMessage(0, Seq(Row(1), Row(2))))
+    writer.add(MemoryWriterCommitMessage(1, Seq(Row(3), Row(4))))
+    writer.add(MemoryWriterCommitMessage(2, Seq(Row(6), Row(7))))
     writer.commit(0)
     assert(sink.latestBatchId.contains(0))
     assert(sink.latestBatchData.map(_.getInt(0)).sorted == Seq(1, 2, 3, 4, 6, 7))
 
-    val newMessages = Seq(
-      MemoryWriterCommitMessage(3, Seq(Row(11), Row(22))),
-      MemoryWriterCommitMessage(0, Seq(Row(33)))
-    )
-    newMessages.foreach(writer.add(_))
+    writer.add(MemoryWriterCommitMessage(3, Seq(Row(11), Row(22))))
+    writer.add(MemoryWriterCommitMessage(0, Seq(Row(33))))
     writer.commit(19)
     assert(sink.latestBatchId.contains(19))
     assert(sink.latestBatchData.map(_.getInt(0)).sorted == Seq(11, 22, 33))
@@ -66,22 +60,16 @@ class MemorySinkV2Suite extends StreamTest with BeforeAndAfter {
   test("microbatch writer") {
     val sink = new MemorySinkV2
     val writer = new MemoryWriter(sink, 0, OutputMode.Append())
-    val messages = Seq(
-      MemoryWriterCommitMessage(0, Seq(Row(1), Row(2))),
-      MemoryWriterCommitMessage(1, Seq(Row(3), Row(4))),
-      MemoryWriterCommitMessage(2, Seq(Row(6), Row(7)))
-    )
-    messages.foreach(writer.add(_))
+    writer.add(MemoryWriterCommitMessage(0, Seq(Row(1), Row(2))))
+    writer.add(MemoryWriterCommitMessage(1, Seq(Row(3), Row(4))))
+    writer.add(MemoryWriterCommitMessage(2, Seq(Row(6), Row(7))))
     writer.commit()
     assert(sink.latestBatchId.contains(0))
     assert(sink.latestBatchData.map(_.getInt(0)).sorted == Seq(1, 2, 3, 4, 6, 7))
 
     val newWriter = new MemoryWriter(sink, 19, OutputMode.Append())
-    val newMessages = Seq(
-      MemoryWriterCommitMessage(3, Seq(Row(11), Row(22))),
-      MemoryWriterCommitMessage(0, Seq(Row(33)))
-    )
-    newMessages.foreach(newWriter.add(_))
+    newWriter.add(MemoryWriterCommitMessage(3, Seq(Row(11), Row(22))))
+    newWriter.add(MemoryWriterCommitMessage(0, Seq(Row(33))))
     newWriter.commit()
     assert(sink.latestBatchId.contains(19))
     assert(sink.latestBatchData.map(_.getInt(0)).sorted == Seq(11, 22, 33))

From d198671aa6794e76f606a364b479b3143bec2c19 Mon Sep 17 00:00:00 2001
From: Wang Gengliang <ltnwgl@gmail.com>
Date: Tue, 30 Jan 2018 20:56:48 +0800
Subject: [PATCH 09/10] Revise comment

---
 .../apache/spark/sql/sources/v2/writer/DataSourceWriter.java  | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataSourceWriter.java b/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataSourceWriter.java
index 4096f2e0b9784..6a3e3d358d61a 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataSourceWriter.java
+++ b/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataSourceWriter.java
@@ -51,8 +51,8 @@
  * 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.
  *
- * In general, all these function calls should be thread-safe in driver side and there is no need
- * to implement concurrency control.
+ * All these methods are guaranteed to be called in a single thread.
+ * No concurrency control is needed.
  *
  * Please refer to the documentation of add/commit/abort methods for detailed specifications.
  */

From 540ff0631471a27af23abb7e8c034bad1ba27cbc Mon Sep 17 00:00:00 2001
From: Wang Gengliang <ltnwgl@gmail.com>
Date: Tue, 30 Jan 2018 22:02:28 +0800
Subject: [PATCH 10/10] revise

---
 .../v2/streaming/writer/StreamWriter.java     | 26 +++++++++++--------
 .../sources/v2/writer/DataSourceWriter.java   | 15 ++++++++---
 2 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/sql/core/src/main/java/org/apache/spark/sql/sources/v2/streaming/writer/StreamWriter.java b/sql/core/src/main/java/org/apache/spark/sql/sources/v2/streaming/writer/StreamWriter.java
index 854b0a2d158e5..83b3e459ba80f 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/sources/v2/streaming/writer/StreamWriter.java
+++ b/sql/core/src/main/java/org/apache/spark/sql/sources/v2/streaming/writer/StreamWriter.java
@@ -32,12 +32,14 @@
 @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()}.
+   * 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
@@ -48,15 +50,17 @@ public interface StreamWriter extends DataSourceWriter {
   /**
    * 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
+   * 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);
 
diff --git a/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataSourceWriter.java b/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataSourceWriter.java
index 6a3e3d358d61a..c4aa9719bcb26 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataSourceWriter.java
+++ b/sql/core/src/main/java/org/apache/spark/sql/sources/v2/writer/DataSourceWriter.java
@@ -51,7 +51,7 @@
  * 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.
  *
- * All these methods are guaranteed to be called in a single thread.
+ * 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.
@@ -68,7 +68,7 @@ public interface DataSourceWriter {
   DataWriterFactory<Row> createWriterFactory();
 
   /**
-   * Handles a commit message which is collected from a successful data writer in the executor side.
+   * 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()}.
@@ -81,6 +81,8 @@ public interface DataSourceWriter {
 
   /**
    * 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.
    *
    * 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
@@ -91,10 +93,17 @@ public interface DataSourceWriter {
   /**
    * 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.
+   * 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 {@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();
 }