HeartSaVioR commented on code in PR #59221:
URL: https://github.com/apache/spark/pull/59221#discussion_r4189852879
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala:
##########
@@ -899,6 +906,59 @@ class MicroBatchExecution(
}
}
+ private def recoverStatefulShufflePartitions(
+ metadata: OffsetSeqMetadataBase,
+ sparkSessionToRunBatches: SparkSession,
+ latestBatchId: Long,
+ committedBatchId: Long): OffsetSeqMetadataBase = {
+ if (metadata.version != OffsetSeqLog.VERSION_2 || committedBatchId < 0) {
+ metadata
+ } else {
+ val hasStatefulShufflePartitions = OffsetSeqMetadata.readValueOpt(
+ metadata, SQLConf.STATEFUL_SHUFFLE_PARTITIONS_INTERNAL).isDefined
+ // In old checkpoints, this config may first appear in an offset written
by a failed restart.
+ // Re-execution commits that offset without rewriting it, so a missing
value in the previous
+ // offset also triggers recovery on the following restart
+ val previousOffsetMissingStatefulShufflePartitions = latestBatchId > 0 &&
+ offsetLog.get(latestBatchId - 1).flatMap(_.metadataOpt)
+ .exists(previous =>
!previous.conf.contains(SQLConf.SHUFFLE_PARTITIONS.key))
+
+ if (hasStatefulShufflePartitions &&
!previousOffsetMissingStatefulShufflePartitions) {
+ metadata
+ } else {
+ metadata match {
+ case v2: OffsetSeqMetadataV2 =>
+ val stateCheckpointLocation = new
Path(checkpointFile("state")).getParent
+ try {
+ val stateMetadataReader = new StateMetadataPartitionReader(
+ stateCheckpointLocation.toString,
+ new SerializableConfiguration(
+ sparkSessionToRunBatches.sessionState.newHadoopConf()),
+ committedBatchId)
+ stateMetadataReader.stateStoreNumPartitions.map { numPartitions
=>
+ logWarning(log"Recovered state-store partition count " +
+ log"${MDC(CONFIG, SQLConf.SHUFFLE_PARTITIONS.key)}=" +
+ log"${MDC(NUM_PARTITIONS, numPartitions)} " +
+ log"from checkpoint state metadata")
+ OffsetSeqMetadata.withStatefulShufflePartitions(v2,
numPartitions)
+ }.getOrElse(metadata)
Review Comment:
Just another case of triggering the same thing:
StateMetadataPartitionReader generally allows the state metadata file to be
absent, but this recovery path requires it to determine the original
state-store partition count. Allowing the missing file to pass through falls
back to the session configuration and recreates the same correctness risk.
We should fail the query and add a regression test covering missing 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]
