ivasilyev1 commented on code in PR #59221:
URL: https://github.com/apache/spark/pull/59221#discussion_r4212925059
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/OffsetSeqLogSuite.scala:
##########
@@ -292,6 +309,242 @@ class OffsetSeqLogSuite extends SharedSparkSession {
}
}
+ private def startStatefulQuery(
+ aggregated: DataFrame,
+ checkpointDir: File,
+ shufflePartitions: Int,
+ queryName: String): StreamingQueryWrapper = {
+ withSQLConf(
+ SQLConf.STREAMING_OFFSET_LOG_FORMAT_VERSION.key -> "2",
+ SQLConf.SHUFFLE_PARTITIONS.key -> shufflePartitions.toString) {
+ aggregated.writeStream
+ .format("memory")
+ .queryName(queryName)
+ .outputMode("complete")
+ .option("checkpointLocation", checkpointDir.getAbsolutePath)
+ .start()
+ .asInstanceOf[StreamingQueryWrapper]
+ }
+ }
+
+ private def createStatefulCheckpoint(
+ checkpointDir: File,
+ queryName: String): (MemoryStream[(Int, Int)], DataFrame) = {
+ val inputData = MemoryStream[(Int, Int)]
+ val aggregated = inputData.toDS().groupBy("_1").count()
+ val query = startStatefulQuery(aggregated, checkpointDir, 5, queryName)
+ try {
+ inputData.addData((1, 0), (2, 0))
+ query.processAllAvailable()
+ inputData.addData((3, 0))
+ query.processAllAvailable()
+ assert(query.streamingQuery.lastExecution.numStateStores === 5)
+ } finally {
+ query.stop()
+ }
+ (inputData, aggregated)
+ }
+
+ private def assertStatefulPartitionRecoveryFails(
+ aggregated: DataFrame,
+ checkpointDir: File,
+ queryName: String): Unit = {
+ val query = startStatefulQuery(aggregated, checkpointDir, 10, queryName)
+ try {
+ val exception = intercept[StreamingQueryException] {
+ query.processAllAvailable()
+ }
+ val message = exception.getCause.getMessage
+ assert(message.contains(
+ "Failed to recover the state-store partition count from checkpoint
metadata"))
+ assert(message.contains("Delete the checkpoint and restart the query"))
+ } finally {
+ query.stop()
+ }
+ }
+
+ private def updateShufflePartitionsInOffset(
+ offsetLog: OffsetSeqLog,
+ checkpointDir: File,
+ batchId: Long,
+ numPartitions: Option[Int]): Unit = {
+ val offsetMap = offsetLog.get(batchId).get match {
+ case offset: OffsetMap => offset
+ case offset => fail(s"Expected a v2 offset, but found
${offset.getClass.getSimpleName}")
+ }
+ val conf = numPartitions match {
+ case Some(value) => offsetMap.metadata.conf.updated(
+ SQLConf.SHUFFLE_PARTITIONS.key, value.toString)
+ case None => offsetMap.metadata.conf - SQLConf.SHUFFLE_PARTITIONS.key
+ }
+ val updatedOffset = offsetMap.copy(metadata = offsetMap.metadata.copy(conf
= conf))
+
+ val offsetFile = new Path(
+ new Path(checkpointDir.getAbsolutePath, "offsets"), batchId.toString)
+ val fileSystem =
offsetFile.getFileSystem(spark.sessionState.newHadoopConf())
+ assert(fileSystem.delete(offsetFile, false))
+ assert(offsetLog.add(batchId, updatedOffset))
+ assert(offsetLog.get(batchId).get.metadataOpt.get.conf
+ .get(SQLConf.SHUFFLE_PARTITIONS.key) === numPartitions.map(_.toString))
+ }
+
+ test("[SPARK-59919] VERSION_2 uses offset partitions when state metadata
differs") {
+ withTempDir { checkpointDir =>
+ val (inputData, aggregated) = createStatefulCheckpoint(
+ checkpointDir, "offsetlog_v2_stateful_restart_test")
+
+ val offsetLog = new OffsetSeqLog(spark,
s"${checkpointDir.getAbsolutePath}/offsets")
+ val metadata = offsetLog.getLatest().get._2.metadataOpt.get
+ assert(metadata.version === OffsetSeqLog.VERSION_2)
+ assert(metadata.conf.get(SQLConf.SHUFFLE_PARTITIONS.key).contains("5"))
+ val latestBatchId = offsetLog.getLatestBatchId().get
+ assert(latestBatchId > 0L)
+ updateShufflePartitionsInOffset(offsetLog, checkpointDir, latestBatchId
- 1, None)
+
+ // Keep the state metadata at 5 and change the latest offset to 10. The
resulting state-store
+ // load failure proves the latest offset value was used instead of state
metadata.
+ updateShufflePartitionsInOffset(offsetLog, checkpointDir, latestBatchId,
Some(10))
+ assert(offsetLog.getLatest().get._2.metadataOpt.get.conf
+ .get(SQLConf.SHUFFLE_PARTITIONS.key).contains("10"))
+ val stateMetadataReader = new StateMetadataPartitionReader(
+ checkpointDir.getAbsolutePath,
+ new SerializableConfiguration(spark.sessionState.newHadoopConf()),
+ latestBatchId)
+ assert(stateMetadataReader.stateStoreNumPartitions.contains(5))
+
+ val query2 = startStatefulQuery(
+ aggregated, checkpointDir, shufflePartitions = 10,
+ queryName = "offsetlog_v2_stateful_restart_test")
+ try {
+ inputData.addData((2, 0), (3, 0))
+ val exception = intercept[StreamingQueryException] {
+ query2.processAllAvailable()
+ }
+ val stateStoreException =
exception.getCause.asInstanceOf[SparkException]
+ assert(stateStoreException.getCondition ===
+ "CANNOT_LOAD_STATE_STORE.CANNOT_READ_DELTA_FILE_NOT_EXISTS")
+ assert(stateStoreException.getMessage.contains("does not exist"))
+ } finally {
+ query2.stop()
+ }
+ }
+ }
+
+ test("[SPARK-59919] VERSION_2 recovers state-store partitions when the
commit log is missing") {
Review Comment:
Switched this + the other stateful tests to use v1 format to be explicit and
not fail unexpectedly in the future
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
