diff --git a/docs/docs/spark/structured-streaming.md b/docs/docs/spark/structured-streaming.md index bee6d80aa331..cf5a970f02ed 100644 --- a/docs/docs/spark/structured-streaming.md +++ b/docs/docs/spark/structured-streaming.md @@ -56,6 +56,44 @@ val stream = df Streaming write also supports [Write merge schema](./sql-write#write-merge-schema). +### 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. + +The commit user is derived from the checkpoint location of the query, so a query keeps it as long +as it keeps its checkpoint. If the checkpoint location never reaches the sink (for example when it +comes from `spark.sql.streaming.checkpointLocation`), the query id Spark stores in the checkpoint +is used instead. Set `write.stream.commit-user` to pin it explicitly, either as an option of the +writer or as a `spark.paimon.write.stream.commit-user` session conf, which is only needed if a +query has to keep its identity across a change of checkpoint location: + +```scala +val stream = df + .writeStream + .outputMode("append") + .option("checkpointLocation", "/path/to/checkpoint") + .option("write.stream.commit-user", "my-streaming-job") + .format("paimon") + .start("/path/to/paimon/sink/table") +``` + +:::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. + +Starting a query from a new checkpoint location gives it a new commit user, so a micro-batch +committed by the previous run is not recognised and its data is written again. + +A table using postpone bucket with `postpone.batch-write-fixed-bucket` commits through a staged +committer that cannot skip a replay; a warning is logged for every such micro-batch. + +::: + ## Streaming Query :::info diff --git a/docs/generated/spark_connector_configuration.html b/docs/generated/spark_connector_configuration.html index d80d14f258ec..f94f0a34be6f 100644 --- a/docs/generated/spark_connector_configuration.html +++ b/docs/generated/spark_connector_configuration.html @@ -140,6 +140,12 @@ Boolean Only effective when 'write.merge-schema' is true. If true, widen an existing column type when the incoming data has a wider compatible type (e.g. INT -> BIGINT, DECIMAL precision increase). Lossy changes are still rejected unless 'write.merge-schema.explicit-cast' is also true. + +
write.stream.commit-user
+ (none) + String + The commit user of a Structured Streaming write. Paimon skips a micro-batch that a previous run of the same query already committed under this user, which is what makes a replayed micro-batch idempotent. By default it is derived from the checkpoint location of the query, so it is stable across restarts; set it explicitly only if the same query has to keep its identity across a change of checkpoint location. +
write.use-v2-write
false diff --git a/paimon-core/src/main/java/org/apache/paimon/table/sink/BatchWriteBuilderImpl.java b/paimon-core/src/main/java/org/apache/paimon/table/sink/BatchWriteBuilderImpl.java index d8c97405e2b0..9f72e702a313 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/sink/BatchWriteBuilderImpl.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/sink/BatchWriteBuilderImpl.java @@ -36,7 +36,8 @@ public class BatchWriteBuilderImpl implements BatchWriteBuilder { private static final long serialVersionUID = 1L; private final InnerTable table; - private final String commitUser; + + private String commitUser; private Map staticPartition; private @Nullable Long rowIdCheckFromSnapshot = null; @@ -61,6 +62,19 @@ public Optional newWriteSelector() { return table.newWriteSelector(); } + /** + * Use a caller-provided commit user instead of the random one. + * + *

A batch job has no reason to do this, but an engine which replays a failed batch with a + * stable identifier (for example a Spark Structured Streaming micro-batch) needs a commit user + * that survives the replay, so that {@link StreamTableCommit#filterAndCommit} can recognise + * what has already been committed. + */ + public BatchWriteBuilderImpl withCommitUser(String commitUser) { + this.commitUser = commitUser; + return this; + } + @Override public BatchWriteBuilder withOverwrite(@Nullable Map staticPartition) { this.staticPartition = staticPartition; @@ -73,7 +87,7 @@ public BatchTableWrite newWrite() { } @Override - public BatchTableCommit newCommit() { + public InnerTableCommit newCommit() { InnerTableCommit commit = table.newCommit(commitUser) .withOverwrite(staticPartition) diff --git a/paimon-core/src/main/java/org/apache/paimon/table/sink/InnerTableCommit.java b/paimon-core/src/main/java/org/apache/paimon/table/sink/InnerTableCommit.java index 43f98d0e7933..d513dd27baec 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/sink/InnerTableCommit.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/sink/InnerTableCommit.java @@ -54,6 +54,17 @@ public interface InnerTableCommit extends StreamTableCommit, BatchTableCommit { InnerTableCommit expireForEmptyCommit(boolean expireForEmptyCommit); + /** + * If this is set to true, {@link StreamTableCommit#filterAndCommit} verifies that every file it + * is about to commit still exists. By default it does. + * + *

The check guards a committable that was restored from an engine's state and may reference + * files deleted long ago. A caller which filters a committable it has just produced itself + * knows those files exist, and can skip a file listing proportional to the size of the + * committable. + */ + InnerTableCommit checkFilesExistence(boolean checkFilesExistence); + InnerTableCommit appendCommitCheckConflict(boolean appendCommitCheckConflict); InnerTableCommit rowIdCheckConflict(@Nullable Long rowIdCheckFromSnapshot); diff --git a/paimon-core/src/main/java/org/apache/paimon/table/sink/TableCommitImpl.java b/paimon-core/src/main/java/org/apache/paimon/table/sink/TableCommitImpl.java index 014b5e64daa1..48cb61eecc78 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/sink/TableCommitImpl.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/sink/TableCommitImpl.java @@ -94,6 +94,7 @@ public class TableCommitImpl implements InnerTableCommit { @Nullable private List overwriteStaticPartitions = null; private boolean batchCommitted = false; private boolean expireForEmptyCommit = true; + private boolean checkFilesExistence = true; public TableCommitImpl( FileStoreCommit commit, @@ -169,6 +170,12 @@ public TableCommitImpl expireForEmptyCommit(boolean expireForEmptyCommit) { return this; } + @Override + public TableCommitImpl checkFilesExistence(boolean checkFilesExistence) { + this.checkFilesExistence = checkFilesExistence; + return this; + } + @Override public TableCommitImpl appendCommitCheckConflict(boolean appendCommitCheckConflict) { commit.appendCommitCheckConflict(appendCommitCheckConflict); @@ -333,13 +340,15 @@ public int filterAndCommitMultiple( List retryCommittables = commit.filterCommitted(sortedCommittables); if (!retryCommittables.isEmpty()) { - checkFilesExistence(retryCommittables); + if (checkFilesExistence) { + verifyFilesExist(retryCommittables); + } commitMultiple(retryCommittables, checkAppendFiles); } return retryCommittables.size(); } - private void checkFilesExistence(List committables) { + private void verifyFilesExist(List committables) { List files = new ArrayList<>(); DataFilePathFactories factories = new DataFilePathFactories(commit.pathFactory()); IndexFilePathFactories indexFactories = new IndexFilePathFactories(commit.pathFactory()); diff --git a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkConnectorOptions.java b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkConnectorOptions.java index 4dd9329d1c4c..c3f8f804e799 100644 --- a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkConnectorOptions.java +++ b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkConnectorOptions.java @@ -109,6 +109,19 @@ public class SparkConnectorOptions { "Wait time in milliseconds between retry attempts for Spark V1 UPDATE " + "on data-evolution tables after row-id range update conflicts."); + public static final ConfigOption STREAM_WRITE_COMMIT_USER = + key("write.stream.commit-user") + .stringType() + .noDefaultValue() + .withDescription( + "The commit user of a Structured Streaming write. Paimon skips a " + + "micro-batch that a previous run of the same query already " + + "committed under this user, which is what makes a replayed " + + "micro-batch idempotent. By default it is derived from the " + + "checkpoint location of the query, so it is stable across " + + "restarts; set it explicitly only if the same query has to " + + "keep its identity across a change of checkpoint location."); + public static final ConfigOption MAX_FILES_PER_TRIGGER = key("read.stream.maxFilesPerTrigger") .intType() diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala index 91efc3d541b1..f91916349217 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala @@ -44,12 +44,13 @@ 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._ import java.io.IOException -import java.util.{Map => JMap} +import java.util.{Collections, Map => JMap} import java.util.Collections.singletonMap import scala.collection.JavaConverters._ @@ -57,8 +58,10 @@ import scala.collection.JavaConverters._ case class PaimonSparkWriter( table: FileStoreTable, writeRowTracking: Boolean = false, - batchId: Option[Long] = None) - extends WriteHelper { + batchId: Option[Long] = None, + commitUser: Option[String] = None) + extends WriteHelper + with Logging { private lazy val tableSchema = table.schema @@ -99,7 +102,13 @@ case class PaimonSparkWriter( if (bucketNum.isPresent) Some(bucketNum.get().intValue()) else None } - val writeBuilder: BatchWriteBuilder = table.newBatchWriteBuilder() + val writeBuilder: BatchWriteBuilder = { + val builder = table.newBatchWriteBuilder() + // A streaming write commits under a commit user that survives a restart, so that a replayed + // micro-batch can be recognised as already committed. + commitUser.foreach(builder.asInstanceOf[BatchWriteBuilderImpl].withCommitUser) + builder + } def withOverwrite(): PaimonSparkWriter = withOverwrite(java.util.Collections.emptyMap()) @@ -454,6 +463,16 @@ case class PaimonSparkWriter( writeBuilder.asInstanceOf[BatchWriteBuilderImpl].rowIdCheckConflict(rowIdCheckFromSnapshot) } + /** + * The commit identifier to deduplicate on, present only for a streaming write that has both a + * batch id and a commit user that is stable across restarts. + */ + private def idempotentCommitIdentifier: Option[Long] = + for { + identifier <- batchId + _ <- commitUser + } yield identifier + def commit(commitMessages: Seq[CommitMessage]): Unit = { commit(commitMessages, null) } @@ -463,6 +482,13 @@ case class PaimonSparkWriter( if (stagedSparkSession == null) { throw new IllegalStateException("Postpone staged write has no SparkSession.") } + idempotentCommitIdentifier.foreach { + identifier => + logWarning( + s"Micro-batch $identifier is written to a postpone bucket table through a staged " + + "commit, which cannot deduplicate a replayed batch. A failure of this query may " + + "duplicate the batch.") + } val finalOperation = Option(operation).getOrElse(Snapshot.Operation.WRITE) val finalMessages = new SparkPostponeStagedCommitter( table, @@ -472,14 +498,29 @@ case class PaimonSparkWriter( postCommit(finalMessages) return } - val activeWriteBuilder = - Option(directPostponeWriteBuilder).getOrElse(writeBuilder) - val tableCommit = activeWriteBuilder.newCommit() + val tableCommit: InnerTableCommit = + if (directPostponeWriteBuilder != null) { + directPostponeWriteBuilder.newCommit() + } else { + writeBuilder.asInstanceOf[BatchWriteBuilderImpl].newCommit() + } if (operation != null) { tableCommit.withOperation(operation) } try { - tableCommit.commit(commitMessages.toList.asJava) + idempotentCommitIdentifier match { + case Some(identifier) => + // Structured Streaming replays a micro-batch with its original batch id after a failure. + // Committing under a stable commit user lets Paimon skip a replay it already committed, + // instead of duplicating the whole batch. The files being committed were written by this + // very batch, so there is no need to list them to prove that they still exist. + tableCommit + .checkFilesExistence(false) + .filterAndCommit( + Collections.singletonMap(Long.box(identifier), commitMessages.toList.asJava)) + case None => + tableCommit.commit(commitMessages.toList.asJava) + } } catch { case e: Throwable => throw new RuntimeException(e); } finally { diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/WriteIntoPaimonTable.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/WriteIntoPaimonTable.scala index 937a47c526e9..2a803a761c1d 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/WriteIntoPaimonTable.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/WriteIntoPaimonTable.scala @@ -40,7 +40,8 @@ case class WriteIntoPaimonTable( saveMode: SaveMode, _data: DataFrame, options: Options, - batchId: Option[Long] = None) + batchId: Option[Long] = None, + commitUser: Option[String] = None) extends RunnableCommand with ExpressionHelper with SchemaEvolutionHelper @@ -58,7 +59,7 @@ case class WriteIntoPaimonTable( updateTableWithOptions( Map(DYNAMIC_PARTITION_OVERWRITE.key -> dynamicPartitionOverwriteMode.toString)) - val writer = PaimonSparkWriter(table, batchId = batchId) + val writer = PaimonSparkWriter(table, batchId = batchId, commitUser = commitUser) if (overwritePartition != null) { writer.withOverwrite(overwritePartition.asJava) } diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/PaimonSink.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/PaimonSink.scala index 9d0a1795b589..6bfec18c4ac1 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/PaimonSink.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/sources/PaimonSink.scala @@ -19,15 +19,21 @@ package org.apache.paimon.spark.sources import org.apache.paimon.options.Options -import org.apache.paimon.spark.{InsertInto, Overwrite} +import org.apache.paimon.spark.{InsertInto, Overwrite, SparkConnectorOptions} import org.apache.paimon.spark.commands.{SchemaEvolutionHelper, WriteIntoPaimonTable} import org.apache.paimon.table.FileStoreTable +import org.apache.spark.internal.Logging import org.apache.spark.sql.{DataFrame, PaimonUtils, SQLContext} import org.apache.spark.sql.execution.streaming.Sink import org.apache.spark.sql.sources.AlwaysTrue import org.apache.spark.sql.streaming.OutputMode +import java.nio.charset.StandardCharsets.UTF_8 +import java.util.UUID + +import scala.collection.JavaConverters._ + class PaimonSink( sqlContext: SQLContext, override val originTable: FileStoreTable, @@ -35,7 +41,66 @@ class PaimonSink( outputMode: OutputMode, options: Options) extends Sink - with SchemaEvolutionHelper { + with SchemaEvolutionHelper + with Logging { + + /** + * Structured Streaming replays a micro-batch with its original batch id when a query is restarted + * after failing between this sink returning from [[addBatch]] and Spark recording the batch as + * completed. Committing every batch under a commit user that is stable across restarts lets + * Paimon skip such a replay instead of committing its data twice. + * + * Resolved lazily: neither the checkpoint location nor the query id is available on the thread + * that constructs the sink. + */ + private lazy val commitUser: String = { + configuredCommitUser.getOrElse { + checkpointLocation + .map(derivedCommitUser("checkpoint", _)) + .orElse(queryId.map(derivedCommitUser("query", _))) + .getOrElse { + logWarning( + "This streaming write has neither a checkpoint location nor a query id to derive a " + + "stable commit user from, so a replayed micro-batch cannot be recognised and may " + + s"be committed twice. Set '${SparkConnectorOptions.STREAM_WRITE_COMMIT_USER.key}' " + + "to make the write idempotent.") + UUID.randomUUID().toString + } + } + } + + // Like the read side, which takes its 'read.stream.*' options from the table, so that a + // 'spark.paimon.' session conf works the same as an option of the writer. + private def configuredCommitUser: Option[String] = { + val fromWriter = options.get(SparkConnectorOptions.STREAM_WRITE_COMMIT_USER) + val fromTable = + Options.fromMap(originTable.options()).get(SparkConnectorOptions.STREAM_WRITE_COMMIT_USER) + Seq(fromWriter, fromTable).find(user => user != null && user.nonEmpty) + } + + // Spark hands the sink its options case-insensitively, but keeps whatever case the user wrote. + private def checkpointLocation: Option[String] = + options.toMap.asScala.collectFirst { + case (key, value) + if key.equalsIgnoreCase(PaimonSink.CHECKPOINT_LOCATION) && value != null && + value.nonEmpty => + value + } + + /** + * The id Spark persists in the checkpoint metadata, hence stable across restarts of the same + * query. It covers the case of a checkpoint location that never reaches the sink options, for + * example one taken from `spark.sql.streaming.checkpointLocation`. It is a thread local of the + * stream execution thread, so it can only be read from within [[addBatch]]. + */ + private def queryId: Option[String] = + Option(sqlContext.sparkContext.getLocalProperty(PaimonSink.QUERY_ID_KEY)).filter(_.nonEmpty) + + private def derivedCommitUser(kind: String, value: String): String = { + val user = s"spark-$kind-${UUID.nameUUIDFromBytes(value.getBytes(UTF_8))}" + logInfo(s"Streaming writes to ${originTable.name()} commit as '$user'.") + user + } override def addBatch(batchId: Long, data: DataFrame): Unit = { val saveMode = if (outputMode == OutputMode.Complete()) { @@ -44,7 +109,18 @@ class PaimonSink( InsertInto } val newData = PaimonUtils.createNewDataFrame(data) - WriteIntoPaimonTable(originTable, saveMode, newData, options, Some(batchId)).run( - sqlContext.sparkSession) + WriteIntoPaimonTable(originTable, saveMode, newData, options, Some(batchId), Some(commitUser)) + .run(sqlContext.sparkSession) } } + +object PaimonSink { + + private val CHECKPOINT_LOCATION = "checkpointLocation" + + /** + * `org.apache.spark.sql.execution.streaming.StreamExecution.QUERY_ID_KEY`, inlined because that + * class is not in the same package across all supported Spark versions. + */ + private val QUERY_ID_KEY = "sql.streaming.queryId" +} diff --git a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/PaimonSinkIdempotencyTest.scala b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/PaimonSinkIdempotencyTest.scala new file mode 100644 index 000000000000..ea0543141a4d --- /dev/null +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/PaimonSinkIdempotencyTest.scala @@ -0,0 +1,323 @@ +/* + * 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.paimon.spark + +import org.apache.paimon.options.Options +import org.apache.paimon.spark.sources.PaimonSink + +import org.apache.spark.SparkConf +import org.apache.spark.sql.{DataFrame, Row} +import org.apache.spark.sql.paimon.shims.memstream.MemoryStream +import org.apache.spark.sql.streaming.{OutputMode, StreamingQuery, StreamTest} + +import java.io.File +import java.util.Collections + +/** + * Structured Streaming guarantees exactly-once only if the sink is idempotent for a repeated + * batchId: when a query fails between the sink returning from `addBatch` and Spark recording the + * batch as completed, the restarted query replays that micro-batch with its original batchId. + */ +class PaimonSinkIdempotencyTest extends PaimonSparkTestBase with StreamTest { + + override protected def sparkConf: SparkConf = { + super.sparkConf.set("spark.sql.catalog.paimon.cache-enabled", "false") + } + + import testImplicits._ + + private def snapshotCount(tableName: String): Long = + loadTable(tableName).snapshotManager().snapshotCount() + + private def latestCommitUser(tableName: String): String = + loadTable(tableName).snapshotManager().latestSnapshot().commitUser() + + private def runToCompletion(query: StreamingQuery): Unit = { + try { + query.processAllAvailable() + } finally { + query.stop() + } + } + + /** + * Leave the checkpoint in the state a driver failure leaves behind when it dies after the sink + * returned from `addBatch` but before Spark recorded the batch: the offset log still has the + * batch, the commit log does not. The restarted query replays it with the same batchId. + */ + private def dropCommitLogEntry(checkpointPath: String, batchId: Long): Unit = { + val commitsDir = new File(checkpointPath, "commits") + val names = Set(batchId.toString, s".$batchId.crc") + val entries = commitsDir.listFiles().filter(f => names.contains(f.getName)) + assert( + entries.exists(_.getName == batchId.toString), + s"no commit log entry for batch $batchId in $commitsDir") + entries.foreach(f => assert(f.delete())) + } + + test("Paimon Sink: replayed micro-batch must not be committed twice") { + failAfter(streamingTimeout) { + withTempDir { + checkpointDir => + spark.sql("CREATE TABLE T (a INT, b STRING)") + val location = loadTable("T").location().toString + val checkpointPath = checkpointDir.getCanonicalPath + + val inputData = MemoryStream[(Int, String)] + val df = inputData.toDS().toDF("a", "b") + inputData.addData((1, "a"), (2, "b"), (3, "c")) + + def start(): StreamingQuery = + df.writeStream + .option("checkpointLocation", checkpointPath) + .format("paimon") + .start(location) + + runToCompletion(start()) + + val expected = Row(1, "a") :: Row(2, "b") :: Row(3, "c") :: Nil + checkAnswer(spark.sql("SELECT * FROM T ORDER BY a"), expected) + assert(snapshotCount("T") == 1) + assert( + latestCommitUser("T").startsWith("spark-checkpoint-"), + s"expected a commit user derived from the checkpoint location, " + + s"but got '${latestCommitUser("T")}'" + ) + + dropCommitLogEntry(checkpointPath, 0) + runToCompletion(start()) + + // The replayed batch must be recognised as already committed. + checkAnswer(spark.sql("SELECT * FROM T ORDER BY a"), expected) + assert( + snapshotCount("T") == 1, + s"replaying batch 0 created a second snapshot (${snapshotCount("T")} in total)") + } + } + } + + test("Paimon Sink: replay is recognised when only the query id is available") { + failAfter(streamingTimeout) { + withTempDir { + checkpointRoot => + spark.sql("CREATE TABLE T (a INT, b STRING)") + val location = loadTable("T").location().toString + val queryName = "paimon_idempotency" + // The location never reaches the sink options this way, so the commit user has to come + // from the query id that Spark persists in the checkpoint metadata. + val checkpointPath = new File(checkpointRoot, queryName).getCanonicalPath + + withSQLConf("spark.sql.streaming.checkpointLocation" -> checkpointRoot.getCanonicalPath) { + val inputData = MemoryStream[(Int, String)] + val df = inputData.toDS().toDF("a", "b") + inputData.addData((1, "a"), (2, "b"), (3, "c")) + + def start(): StreamingQuery = + df.writeStream + .queryName(queryName) + .format("paimon") + .start(location) + + runToCompletion(start()) + + val expected = Row(1, "a") :: Row(2, "b") :: Row(3, "c") :: Nil + checkAnswer(spark.sql("SELECT * FROM T ORDER BY a"), expected) + assert(snapshotCount("T") == 1) + assert( + latestCommitUser("T").startsWith("spark-query-"), + s"expected a commit user derived from the query id, " + + s"but got '${latestCommitUser("T")}'") + + dropCommitLogEntry(checkpointPath, 0) + runToCompletion(start()) + + checkAnswer(spark.sql("SELECT * FROM T ORDER BY a"), expected) + assert( + snapshotCount("T") == 1, + s"replaying batch 0 created a second snapshot (${snapshotCount("T")} in total)") + } + } + } + } + + test("Paimon Sink: replay of a batch that is not the first one is recognised") { + failAfter(streamingTimeout) { + withTempDir { + checkpointDir => + spark.sql("CREATE TABLE T (a INT, b STRING)") + val location = loadTable("T").location().toString + val checkpointPath = checkpointDir.getCanonicalPath + + val inputData = MemoryStream[(Int, String)] + val df = inputData.toDS().toDF("a", "b") + + def start(): StreamingQuery = + df.writeStream + .option("checkpointLocation", checkpointPath) + .format("paimon") + .start(location) + + val query = start() + try { + inputData.addData((1, "a")) + query.processAllAvailable() + inputData.addData((2, "b")) + query.processAllAvailable() + inputData.addData((3, "c")) + query.processAllAvailable() + } finally { + query.stop() + } + + val expected = Row(1, "a") :: Row(2, "b") :: Row(3, "c") :: Nil + checkAnswer(spark.sql("SELECT * FROM T ORDER BY a"), expected) + assert(snapshotCount("T") == 3) + + dropCommitLogEntry(checkpointPath, 2) + runToCompletion(start()) + + checkAnswer(spark.sql("SELECT * FROM T ORDER BY a"), expected) + assert( + snapshotCount("T") == 3, + s"replaying batch 2 created another snapshot (${snapshotCount("T")} in total)") + } + } + } + + test("Paimon Sink: replayed micro-batch of a complete mode query is not committed twice") { + failAfter(streamingTimeout) { + withTempDir { + checkpointDir => + spark.sql("CREATE TABLE T (city STRING, population LONG)") + val location = loadTable("T").location().toString + val checkpointPath = checkpointDir.getCanonicalPath + + val inputData = MemoryStream[(Int, String)] + val df = inputData + .toDS() + .toDF("uid", "city") + .groupBy("city") + .count() + .toDF("city", "population") + inputData.addData((1, "HZ"), (2, "BJ"), (3, "BJ")) + + def start(): StreamingQuery = + df.writeStream + .outputMode("complete") + .option("checkpointLocation", checkpointPath) + .format("paimon") + .start(location) + + runToCompletion(start()) + + val expected = Row("BJ", 2L) :: Row("HZ", 1L) :: Nil + checkAnswer(spark.sql("SELECT * FROM T ORDER BY city"), expected) + val snapshotsAfterFirstBatch = snapshotCount("T") + + dropCommitLogEntry(checkpointPath, 0) + runToCompletion(start()) + + checkAnswer(spark.sql("SELECT * FROM T ORDER BY city"), expected) + assert( + snapshotCount("T") == snapshotsAfterFirstBatch, + s"replaying batch 0 created another snapshot (${snapshotCount("T")} in total, " + + s"$snapshotsAfterFirstBatch before the replay)" + ) + } + } + } + + test("Paimon Sink: write.stream.commit-user overrides the derived commit user") { + failAfter(streamingTimeout) { + withTempDir { + checkpointDir => + spark.sql("CREATE TABLE T (a INT, b STRING)") + val location = loadTable("T").location().toString + val checkpointPath = checkpointDir.getCanonicalPath + + val inputData = MemoryStream[(Int, String)] + val df = inputData.toDS().toDF("a", "b") + inputData.addData((1, "a"), (2, "b"), (3, "c")) + + def start(): StreamingQuery = + df.writeStream + .option("checkpointLocation", checkpointPath) + .option("write.stream.commit-user", "my-streaming-job") + .format("paimon") + .start(location) + + runToCompletion(start()) + + val expected = Row(1, "a") :: Row(2, "b") :: Row(3, "c") :: Nil + checkAnswer(spark.sql("SELECT * FROM T ORDER BY a"), expected) + assert(latestCommitUser("T") == "my-streaming-job") + + dropCommitLogEntry(checkpointPath, 0) + runToCompletion(start()) + + checkAnswer(spark.sql("SELECT * FROM T ORDER BY a"), expected) + assert( + snapshotCount("T") == 1, + s"replaying batch 0 created a second snapshot (${snapshotCount("T")} in total)") + } + } + } + + test("Paimon Sink: write.stream.commit-user can come from a session conf") { + failAfter(streamingTimeout) { + withTempDir { + checkpointDir => + spark.sql("CREATE TABLE T (a INT, b STRING)") + val location = loadTable("T").location().toString + + withSQLConf("spark.paimon.write.stream.commit-user" -> "job-from-conf") { + val inputData = MemoryStream[(Int, String)] + val df = inputData.toDS().toDF("a", "b") + inputData.addData((1, "a")) + + runToCompletion( + df.writeStream + .option("checkpointLocation", checkpointDir.getCanonicalPath) + .format("paimon") + .start(location)) + } + + checkAnswer(spark.sql("SELECT * FROM T"), Row(1, "a") :: Nil) + assert(latestCommitUser("T") == "job-from-conf") + } + } + } + + test("Paimon Sink: addBatch with a repeated batchId must be a no-op") { + spark.sql("CREATE TABLE T2 (a INT, b STRING)") + val sink = new PaimonSink( + spark.sqlContext, + loadTable("T2"), + Nil, + OutputMode.Append(), + Options.fromMap(Collections.singletonMap("write.stream.commit-user", "direct-api"))) + + val batch: DataFrame = Seq((1, "a"), (2, "b")).toDF("a", "b") + sink.addBatch(0L, batch) + sink.addBatch(0L, batch) + + checkAnswer(spark.sql("SELECT * FROM T2 ORDER BY a"), Row(1, "a") :: Row(2, "b") :: Nil) + assert(snapshotCount("T2") == 1) + } +}