HeartSaVioR commented on code in PR #59221:
URL: https://github.com/apache/spark/pull/59221#discussion_r4190287914


##########
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:
   If we see the stateful operator in the query but we don't see the shuffle 
partitions in offset log, and we fall back to find the state metadata and fail 
to read the shuffle partitions for that operator "in any reason", we must fail 
the query. Please focus on the part of "in any reason".



-- 
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]

Reply via email to