ivasilyev1 commented on code in PR #59221:
URL: https://github.com/apache/spark/pull/59221#discussion_r4194134874
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala:
##########
@@ -899,6 +906,75 @@ class MicroBatchExecution(
}
}
+ private def recoverStatefulShufflePartitions(
+ metadata: OffsetSeqMetadataBase,
+ sparkSessionToRunBatches: SparkSession,
+ latestBatchId: Long,
+ committedBatchId: Long): OffsetSeqMetadataBase = {
+ if (metadata.version != OffsetSeqLog.VERSION_2) {
+ 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
+ val stateMetadataBatchId = if (committedBatchId >= 0) {
+ committedBatchId
+ } else {
+ latestBatchId
+ }
+ val numPartitionsOpt = try {
+ val stateMetadataReader = new StateMetadataPartitionReader(
+ stateCheckpointLocation.toString,
+ new SerializableConfiguration(
+ sparkSessionToRunBatches.sessionState.newHadoopConf()),
+ stateMetadataBatchId)
+ val partitionsOpt = stateMetadataReader.stateStoreNumPartitions
+ // If the query was stateful and we could not get the number of
partitions
+ // for any reason, abort the query.
+ if (partitionsOpt.isEmpty &&
containsStatefulOperator(analyzedPlan)) {
+ throw new IllegalStateException(
Review Comment:
This was just an easy way to keep some of the logic a bit more condensed
since this gets immediately caught and re raised with SparkException. I'll
reorganize and remove while addressing the other comment
--
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]
