HeartSaVioR commented on code in PR #59221:
URL: https://github.com/apache/spark/pull/59221#discussion_r4202312856
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala:
##########
@@ -899,6 +906,50 @@ class MicroBatchExecution(
}
}
+ private def recoverStatefulShufflePartitions(
+ metadata: OffsetSeqMetadataBase,
+ sparkSessionToRunBatches: SparkSession,
+ latestBatchId: Long,
+ committedBatchId: Long): OffsetSeqMetadataBase = {
+ if (metadata.version != OffsetSeqLog.VERSION_2 ||
+ OffsetSeqMetadata.readValueOpt(
+ metadata, SQLConf.STATEFUL_SHUFFLE_PARTITIONS_INTERNAL).isDefined ||
+ !containsStatefulOperator(analyzedPlan)) {
+ metadata
+ } else {
+ metadata match {
+ case v2: OffsetSeqMetadataV2 =>
+ val stateCheckpointLocation = new
Path(checkpointFile("state")).getParent
+ val stateMetadataBatchId = if (committedBatchId >= 0)
committedBatchId else latestBatchId
+ val failedRecoveryMessage =
+ s"Failed to recover the state-store partition count from
checkpoint " +
+ s"metadata at $stateCheckpointLocation. This can happen if the
checkpoint " +
+ "was created using offset log format V2 in an unpatched Spark
distribution " +
Review Comment:
nit: an unpatched <- this does not provide any context. If we want to be
honest, we need to refer to the issue explicitly.
I'd suggest changing the sentence of "This can happen if the checkpoint ~
and state metadata was subsequently corrupted." to below:
This can happen if the checkpoint was created using offset log format V2 and
state metadata was subsequently corrupted due to a bug. See SPARK-59919 for
more details.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/OffsetSeqLogSuite.scala:
##########
@@ -47,6 +52,18 @@ class OffsetSeqLogSuite extends SharedSparkSession {
assert(metadata.conf.get(SQLConf.ALLOW_EXCEPT_ON_STREAMING_DATAFRAME.key).contains("false"))
}
+ test("[SPARK-59919] v1 and v2 metadata persist rebound stateful shuffle
partitions") {
Review Comment:
nit: common practice is `SPARK-59919: ` rather than `[SPARK-59919]`
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala:
##########
@@ -825,9 +826,15 @@ class MicroBatchExecution(
// update offset metadata
nextOffsets.metadataOpt.foreach { metadata =>
- OffsetSeqMetadata.setSessionConf(metadata,
sparkSessionToRunBatches.sessionState.conf)
+ val committedBatchId = commitLog.getLatestBatchId().getOrElse(-1L)
+ val metadataWithRecoveredPartitions =
recoverStatefulShufflePartitions(
+ metadata, sparkSessionToRunBatches, latestBatchId,
committedBatchId)
+ OffsetSeqMetadata.setSessionConf(
+ metadataWithRecoveredPartitions,
sparkSessionToRunBatches.sessionState.conf)
execCtx.offsetSeqMetadata = OffsetSeqMetadata(
- metadata.batchWatermarkMs, metadata.batchTimestampMs,
sparkSessionToRunBatches.conf)
+ metadataWithRecoveredPartitions.batchWatermarkMs,
+ metadataWithRecoveredPartitions.batchTimestampMs,
+ sparkSessionToRunBatches.conf)
watermarkTracker = WatermarkTracker(sparkSessionToRunBatches.conf,
logicalPlan)
watermarkTracker.setWatermark(metadata.batchWatermarkMs)
Review Comment:
nit: since we build metadataWithRecoveredPartitions, probably good to be
consistent, use `metadataWithRecoveredPartitions` rather than `metadata`.
##########
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]
Review Comment:
Maybe, could you give below a try and whether below can replace this
validation? This validation is coupled with implementation detail for HDFS
state store provider which is good to be avoided.
```
// The failure is expected because physical state exists for only five
partitions.
// Planning ten stores proves the latest offset value won over state
metadata.
assert(query2.streamingQuery.lastExecution.numStateStores === 10)
```
##########
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:
For this test, we'd be better to explicitly set the checkpoint version to 1.
Eventually we could flip the default and it would fail this test.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala:
##########
@@ -825,9 +826,15 @@ class MicroBatchExecution(
// update offset metadata
nextOffsets.metadataOpt.foreach { metadata =>
- OffsetSeqMetadata.setSessionConf(metadata,
sparkSessionToRunBatches.sessionState.conf)
+ val committedBatchId = commitLog.getLatestBatchId().getOrElse(-1L)
Review Comment:
listing is a costly operation, especially against object store. We do the
listing in L829 (here) and L845 (commitLog.getLatest()) - we can combine them
and only do it once.
```
val latestCommittedBatch = commitLog.getLatest()
val committedBatchId = latestCommittedBatch.map(_._1).getOrElse(-1L)
```
Do this before L827, and remove L829 and replace L845 with
`latestCommittedBatch match {`
##########
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]
Review Comment:
nit: shall we use checkError for the validation if it's available?
##########
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()
+ }
+ }
+ }
Review Comment:
I don't see the test for fallback case, offset log is missing shuffle
partitions (impacted) but the checkpoint has a readable state metadata which
contains the shuffle partitions, and we continue with it. Please make sure we
change the SQL config on restart to make the shuffle partitions in session
config differs from shuffle partitions in state metadata (clearly, value in
session config > value in state metadata, so that the query will fail if it
doesn't read the state metadata).
--
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]
