[spark] Commit a streaming micro-batch through the stream write API

docs/docs/spark/structured-streaming.md

@@ -66,6 +66,58 @@ Use a durable checkpoint location accessible to the cluster for deployed jobs. G
 its own checkpoint directory. Stop this example with `writer.stop()`.
 Streaming writes also support [Schema Evolution on Write](./schema-evolution).
 
+### Exactly-once
+
+Structured Streaming replays a micro-batch with its original batch id when a query is restarted
+after failing between the sink writing the batch and Spark recording that batch as completed.
+Paimon commits every micro-batch under a commit user that is stable across restarts, and skips a
+batch that the same user already committed, so a replay does not write the data twice. A query in
+`complete` mode is the exception: each of its micro-batches overwrites the table with the complete
+result, so its replay is committed again and leaves the same rows, as before. Micro-batch
+`n` is committed under commit identifier `n + 1`, the way Flink numbers its checkpoints, which is
+what the `$snapshots` system table shows and what a `compacted-full` scan recognises a scheduled
+full compaction by.
+
+What the commit user identifies is one incarnation of a checkpoint, since the batch ids it
+numbers are only unique within one: reusing it across two checkpoints would make Paimon skip the
+data of the second one, while changing it within one query would bring the duplicate back. It is
+therefore derived from the query id that Spark persists in the checkpoint, which is new when a
+checkpoint is recreated, unchanged when a query resumes from one, and independent of how the
+location is spelled, and it cannot be set explicitly. `commit.user-prefix` is read from the table as
+usual and prefixes the derived user, so a job keeps the name its table was configured with.
+
+:::note
+
+A skipped replay leaves the data files it wrote behind, uncommitted. They are removed by
+[orphan file cleaning](../maintenance/manage-snapshots#remove-orphan-files), like any other
+uncommitted file.
+
+A query that starts from a new checkpoint gets a new commit user, so a micro-batch the previous
+run committed is not recognised and its data is written again.
+
+A replay is recognised by the snapshots its commit user left behind, which snapshot expiration
+eventually removes. Before committing a micro-batch, the sink records under the checkpoint location
+of the query which snapshot was the latest before its first attempt. Retries preserve this original
+marker rather than advancing it. If filtering cannot recognise the replay and expiration has removed
+the snapshot immediately after that boundary, the query fails instead of possibly writing it twice.
+An incomplete or corrupt marker also fails the query, since it may belong to an already committed
+batch. The error names the marker file; restore it or delete it only after verifying that the batch
+was not committed. Retain snapshots for longer than a query may be down.
+
+A streaming write to a postpone bucket table writes the postpone bucket, as a Flink streaming job
+does, whatever `postpone.batch-write-fixed-bucket` says: that option is for a batch job which ends
+with its commit. The rows become readable once a
+[compaction](../primary-key-table/data-distribution#postpone-bucket) has sorted them into real
+buckets.
+
+:::
+
+The committer of a query is created with its first micro-batch and lives until the query
+terminates, again as a Flink committer lives across checkpoints. Only the first micro-batch of a
+run looks up what a previous run committed; the maintenance a commit schedules, such as tag
+creation and snapshot or partition expiration, runs while the next micro-batch is written when
+`snapshot.expire.execution-mode` is `async`.
+
 ## Streaming Query
 
 :::info

paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala

@@ -37,14 +37,15 @@ import org.apache.paimon.spark.schema.SparkSystemColumns.{BUCKET_COL, ROW_KIND_C
 import org.apache.paimon.spark.sort.TableSorter
 import org.apache.paimon.spark.util.OptionUtils.paimonExtensionEnabled
 import org.apache.paimon.spark.util.SparkRowUtils
-import org.apache.paimon.spark.write.{PaimonDataWrite, WriteHelper, WriteTaskResult}
+import org.apache.paimon.spark.write.{PaimonDataWrite, StreamingWrite, WriteHelper, WriteTaskResult}
 import org.apache.paimon.table.{FileStoreTable, SpecialFields}
 import org.apache.paimon.table.BucketMode._
 import org.apache.paimon.table.sink._
 import org.apache.paimon.types.RowKind
 import org.apache.paimon.utils.{SerializationUtils, UriReaderFactory}
 
 import org.apache.spark.{Partitioner, TaskContext}
+import org.apache.spark.internal.Logging
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql._
 import org.apache.spark.sql.functions._
@@ -55,11 +56,20 @@ import java.util.Collections.singletonMap
 
 import scala.collection.JavaConverters._
 
+/**
+ * @param streaming
+ *   the micro-batch this writer writes, when it writes for a streaming query. Such a write goes
+ *   through the stream write API, the way a Flink streaming job does: it is committed under an
+ *   identifier by the committer of the query, which recognises a replayed micro-batch.
+ */
 case class PaimonSparkWriter(
     table: FileStoreTable,
     writeRowTracking: Boolean = false,
-    batchId: Option[Long] = None)
-  extends WriteHelper {
+    streaming: Option[StreamingWrite] = None)
+  extends WriteHelper
+  with Logging {
+
+  private val commitIdentifier: Option[Long] = streaming.map(_.commitIdentifier)
 
   private lazy val tableSchema = table.schema
 
@@ -89,8 +99,14 @@ case class PaimonSparkWriter(
 
   @transient private lazy val metricRegistry = SparkMetricRegistry()
 
+  /**
+   * A streaming write to a postpone bucket table always goes to bucket -2, to be bucketed by a
+   * later compaction, the same as a Flink streaming write: the fixed-bucket paths are for a batch
+   * job that ends with its commit.
+   */
   val postponeBatchWriteFixedBucket: Boolean =
-    table.bucketMode() == POSTPONE_MODE && coreOptions.postponeBatchWriteFixedBucket()
+    streaming.isEmpty && table.bucketMode() == POSTPONE_MODE &&
+      coreOptions.postponeBatchWriteFixedBucket()
 
   private val postponeBaseSnapshotId =
     if (postponeBatchWriteFixedBucket)
@@ -102,25 +118,34 @@ case class PaimonSparkWriter(
     if (bucketNum.isPresent) Some(bucketNum.get().intValue()) else None
   }
 
-  val writeBuilder: BatchWriteBuilder = table.newBatchWriteBuilder()
+  val writeBuilder: WriteBuilder = streaming match {
+    case Some(streamingWrite) =>
+      // Every micro-batch is committed under the commit user of the query, which is stable across
+      // its runs, so that a replayed micro-batch can be recognised as already committed.
+      table.newStreamWriteBuilder().withCommitUser(streamingWrite.context.commitUser)
+    case None => table.newBatchWriteBuilder()
+  }
 
   def withOverwrite(): PaimonSparkWriter = withOverwrite(java.util.Collections.emptyMap())
 
   def withOverwrite(partition: JMap[String, String]): PaimonSparkWriter = {
     overwritePartitionSpec = Some(partition.asScala.toMap)
-    if (!postponeBatchWriteFixedBucket) {
-      writeBuilder.withOverwrite(partition)
+    writeBuilder match {
+      case batchWriteBuilder: BatchWriteBuilder if !postponeBatchWriteFixedBucket =>
+        batchWriteBuilder.withOverwrite(partition)
+      // A stream write builder is told at commit time; see [[commit]].
+      case _ =>
     }
     this
   }
 
   def writeOnly(): PaimonSparkWriter = {
-    PaimonSparkWriter(table.copy(singletonMap(WRITE_ONLY.key(), "true")))
+    PaimonSparkWriter(table.copy(singletonMap(WRITE_ONLY.key(), "true")), streaming = streaming)
   }
 
   def withRowTracking(): PaimonSparkWriter = {
     if (coreOptions.rowTrackingEnabled()) {
-      PaimonSparkWriter(table, writeRowTracking = true)
+      PaimonSparkWriter(table, writeRowTracking = true, streaming = streaming)
     } else {
       this
     }
@@ -140,7 +165,7 @@ case class PaimonSparkWriter(
       } else {
         None
       }
-    val activeWriteBuilder: BatchWriteBuilder = directPostponeBucketNum match {
+    val activeWriteBuilder: WriteBuilder = directPostponeBucketNum match {
       case Some(_) =>
         val directWriteOptions = new java.util.HashMap[String, String]()
         directWriteOptions.put(
@@ -176,9 +201,11 @@ case class PaimonSparkWriter(
         rowKindColIdx,
         writeRowTracking,
         fullCompactionDeltaCommits,
-        batchId,
+        commitIdentifier,
         uriReaderFactory,
-        postponePartitionBucketComputer
+        postponePartitionBucketComputer,
+        // The stream write builder does not know that the commit overwrites, see [[commit]].
+        ignorePreviousFiles = streaming.isDefined && overwritePartitionSpec.isDefined
       )
 
     def sparkParallelism = {
@@ -458,6 +485,18 @@ case class PaimonSparkWriter(
   }
 
   def commit(commitMessages: Seq[CommitMessage], operation: Snapshot.Operation): Unit = {
+    streaming match {
+      case Some(streamingWrite) =>
+        commitMicroBatch(streamingWrite, commitMessages, operation)
+        postCommit(commitMessages)
+      case None =>
+        commitBatch(commitMessages, operation)
+    }
+  }
+
+  private def commitBatch(
+      commitMessages: Seq[CommitMessage],
+      operation: Snapshot.Operation): Unit = {
     if (postponeBatchWriteFixedBucket && directPostponeWriteBuilder == null) {
       if (stagedSparkSession == null) {
         throw new IllegalStateException("Postpone staged write has no SparkSession.")
@@ -471,8 +510,8 @@ case class PaimonSparkWriter(
       postCommit(finalMessages)
       return
     }
-    val activeWriteBuilder =
-      Option(directPostponeWriteBuilder).getOrElse(writeBuilder)
+    val activeWriteBuilder: BatchWriteBuilder =
+      Option(directPostponeWriteBuilder).getOrElse(writeBuilder.asInstanceOf[BatchWriteBuilder])
     val tableCommit = activeWriteBuilder.newCommit()
     tableCommit.withMetricRegistry(metricRegistry)
     if (operation != null) {
@@ -488,6 +527,38 @@ case class PaimonSparkWriter(
     postCommit(commitMessages)
   }
 
+  /**
+   * Commits a micro-batch through the committer of the query, which the first micro-batch of a run
+   * creates and every later one reuses.
+   */
+  private def commitMicroBatch(
+      streamingWrite: StreamingWrite,
+      commitMessages: Seq[CommitMessage],
+      operation: Snapshot.Operation): Unit = {
+    streamingWrite.context.openCommitter {
+      val committer = writeBuilder
+        .asInstanceOf[StreamWriteBuilder]
+        .newCommit()
+        .asInstanceOf[InnerTableCommit]
+      committer.withMetricRegistry(metricRegistry)
+      // Both are the same for every micro-batch of a query: its output mode does not change, and
+      // a batch that wrote nothing is no reason to publish a snapshot, as for a batch write.
+      overwritePartitionSpec.foreach(spec => committer.withOverwrite(spec.asJava))
+      committer.ignoreEmptyCommit(
+        Option(coreOptions.toConfiguration.get(CoreOptions.SNAPSHOT_IGNORE_EMPTY_COMMIT))
+          .forall(_.booleanValue()))
+      if (operation != null) {
+        committer.withOperation(operation)
+      }
+      committer
+    }
+    streamingWrite.context.commit(
+      streamingWrite,
+      commitMessages,
+      table,
+      overwrite = overwritePartitionSpec.isDefined)
+  }
+
   private def baseSnapshotHasNoRealBuckets: Boolean = {
     postponeBaseSnapshotId.forall {
       snapshotId =>

paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/WriteIntoPaimonTable.scala

@@ -24,7 +24,7 @@ import org.apache.paimon.options.Options
 import org.apache.paimon.spark._
 import org.apache.paimon.spark.catalyst.analysis.ReplacePaimonFunctions
 import org.apache.paimon.spark.catalyst.analysis.expressions.ExpressionHelper
-import org.apache.paimon.spark.write.PaimonWriteOptions
+import org.apache.paimon.spark.write.{PaimonWriteOptions, StreamingWrite}
 import org.apache.paimon.table.FileStoreTable
 
 import org.apache.spark.internal.Logging
@@ -40,7 +40,7 @@ case class WriteIntoPaimonTable(
     saveMode: SaveMode,
     _data: DataFrame,
     options: Options,
-    batchId: Option[Long] = None)
+    streaming: Option[StreamingWrite] = None)
   extends RunnableCommand
   with ExpressionHelper
   with SchemaEvolutionHelper
@@ -58,7 +58,7 @@ case class WriteIntoPaimonTable(
     updateTableWithOptions(
       Map(DYNAMIC_PARTITION_OVERWRITE.key -> dynamicPartitionOverwriteMode.toString))
 
-    val writer = PaimonSparkWriter(table, batchId = batchId)
+    val writer = PaimonSparkWriter(table, streaming = streaming)
     if (overwritePartition != null) {
       writer.withOverwrite(overwritePartition.asJava)
     }