This is an automated email from the ASF dual-hosted git repository.
HeartSaVioR pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 5d17442dec43 [SPARK-59919][SS] Fix OffsetSeqMetadataV2 missing shuffle
partitions on rebinding
5d17442dec43 is described below
commit 5d17442dec43f6dbea12973f9da18197b0e8d27c
Author: ivasilyev1 <[email protected]>
AuthorDate: Thu Oct 8 15:17:35 2026 +0900
[SPARK-59919][SS] Fix OffsetSeqMetadataV2 missing shuffle partitions on
rebinding
### What changes were proposed in this pull request?
The changes in this PR:
- Ensure that both V1 and V2 offset logs emit rebound SQL conf
- Attempt to perform recovery of missing `spark.sql.shuffle.partitions` in
V2 offset log by obtaining value from state metadata if possible
### Why are the changes needed?
With current implementation the stateful shuffle partitions are not written
out by OffsetSeqMetadataV2. If query is restarted and the session configured
shuffle partitions change, the query will fail.
The exact error scenario and steps to reproduce are described in:
https://issues.apache.org/jira/browse/SPARK-59919
The missing config key also leads to possible crash scenario as in
https://github.com/apache/spark/pull/58929, which was addressed separately in
that code path.
### Does this PR introduce _any_ user-facing change?
Yes, see the detail in https://issues.apache.org/jira/browse/SPARK-59919
### How was this patch tested?
Tested with the included unit tests in this PR as well as with manual
reproduction in spark
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: GPT 5.6 Luna
Closes #59221 from ivasilyev1/v2-shuffle-partitions.
Authored-by: ivasilyev1 <[email protected]>
Signed-off-by: Jungtaek Lim <[email protected]>
---
.../v2/state/metadata/StateMetadataSource.scala | 14 +
.../streaming/checkpointing/OffsetSeq.scala | 48 +++-
.../streaming/runtime/MicroBatchExecution.scala | 97 +++++--
.../execution/streaming/OffsetSeqLogSuite.scala | 304 ++++++++++++++++++++-
4 files changed, 426 insertions(+), 37 deletions(-)
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
index cef3a7a7e5cd..3f2dca133fa5 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/metadata/StateMetadataSource.scala
@@ -241,6 +241,20 @@ class StateMetadataPartitionReader(
}
}
+ private[sql] def stateStoreNumPartitions: Option[Int] = {
+ allOperatorStateMetadata
+ .flatMap(_.stateStoresMetadata.map(_.numPartitions))
+ .distinct
+ .toSeq match {
+ case Seq() => None
+ case Seq(numPartitions) => Some(numPartitions)
+ case values =>
+ throw new IllegalStateException(
+ s"Found conflicting state-store partition counts in checkpoint: " +
+ values.mkString("[", ", ", "]"))
+ }
+ }
+
// From v2, we also need to populate the operatorProperties and
stateSchemaFilePath fields
// for use with the state data source reader
private[sql] lazy val stateMetadata: Iterator[StateMetadataTableEntry] = {
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/checkpointing/OffsetSeq.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/checkpointing/OffsetSeq.scala
index e0442b29bfd7..5731b9d172d5 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/checkpointing/OffsetSeq.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/checkpointing/OffsetSeq.scala
@@ -232,6 +232,26 @@ object OffsetSeqMetadata extends Logging {
private val rebindSQLConfsOffsetLogToSession: Map[ConfigEntry[_],
ConfigEntry[_]] =
rebindSQLConfsSessionToOffsetLog.map { case (k, v) => (v, k) }.toMap
+ /**
+ * Builds the configuration map persisted in offset metadata.
+ *
+ * Re-bound session configurations are written using their legacy offset-log
+ * keys for checkpoint compatibility.
+ */
+ private[checkpointing] def confsForOffsetLog(
+ sessionConf: RuntimeConfig): Map[String, String] = {
+ val relevantConfs = relevantSQLConfs.map { conf =>
+ conf.key -> sessionConf.get(conf.key)
+ }.toMap
+
+ val reboundConfs = rebindSQLConfsSessionToOffsetLog.map {
+ case (confInSession, confInOffsetLog) =>
+ confInOffsetLog.key -> sessionConf.get(confInSession.key)
+ }.toMap
+
+ relevantConfs ++ reboundConfs
+ }
+
/**
* Default values of relevant configurations that are used for backward
compatibility.
* As new configurations are added to the metadata, existing checkpoints may
not have those
@@ -283,12 +303,19 @@ object OffsetSeqMetadata extends Logging {
batchWatermarkMs: Long,
batchTimestampMs: Long,
sessionConf: RuntimeConfig): OffsetSeqMetadata = {
- val confs = relevantSQLConfs.map { conf => conf.key ->
sessionConf.get(conf.key) }.toMap
- val confsFromRebind = rebindSQLConfsSessionToOffsetLog.map {
- case (confInSession, confInOffsetLog) =>
- confInOffsetLog.key -> sessionConf.get(confInSession.key)
- }.toMap
- OffsetSeqMetadata(batchWatermarkMs, batchTimestampMs, confs++
confsFromRebind)
+ OffsetSeqMetadata(batchWatermarkMs, batchTimestampMs,
confsForOffsetLog(sessionConf))
+ }
+
+ /**
+ * Adds the stateful shuffle partitions to offset metadata if missing. The
value
+ * is written with the legacy offset-log key so that the regular session
configuration
+ * rebinding path can restore it.
+ */
+ private[sql] def withStatefulShufflePartitions(
+ metadata: OffsetSeqMetadataV2,
+ numPartitions: Int): OffsetSeqMetadataV2 = {
+ val conf = metadata.conf.updated(SHUFFLE_PARTITIONS.key,
numPartitions.toString)
+ metadata.copy(conf = conf)
}
/** Set the SparkSession configuration with the values in the metadata */
@@ -455,9 +482,10 @@ object OffsetSeqMetadataV2 {
batchWatermarkMs: Long,
batchTimestampMs: Long,
sessionConf: RuntimeConfig): OffsetSeqMetadataV2 = {
- val confs = OffsetSeqMetadata.relevantSQLConfs.map {
- conf => conf.key -> sessionConf.get(conf.key)
- }.toMap
- OffsetSeqMetadataV2(batchWatermarkMs, batchTimestampMs, confs)
+ OffsetSeqMetadataV2(
+ batchWatermarkMs,
+ batchTimestampMs,
+ OffsetSeqMetadata.confsForOffsetLog(sessionConf)
+ )
}
}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala
index 4bb6124d37da..1bda5691f80b 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala
@@ -45,8 +45,9 @@ import org.apache.spark.sql.errors.QueryExecutionErrors
import org.apache.spark.sql.execution.{SparkPlan, SQLExecution}
import org.apache.spark.sql.execution.datasources.LogicalRelation
import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation,
RealTimeStreamScanExec, StreamingDataSourceV2Relation,
StreamingDataSourceV2ScanRelation, StreamWriterCommitProgress,
WriteToDataSourceV2Exec}
+import
org.apache.spark.sql.execution.datasources.v2.state.metadata.StateMetadataPartitionReader
import org.apache.spark.sql.execution.streaming.{AvailableNowTrigger, Offset,
OneTimeTrigger, ProcessingTimeTrigger, RealTimeTrigger, Sink, Source,
StreamingQueryPlanTraverseHelper}
-import
org.apache.spark.sql.execution.streaming.checkpointing.{CheckpointFileManager,
CheckpointVersionManager, CommitLog, CommitLogType, CommitMetadataV3,
OffsetLogType, OffsetSeqBase, OffsetSeqLog, OffsetSeqMetadata,
OffsetSeqMetadataV2, SinkMetadataInfo}
+import
org.apache.spark.sql.execution.streaming.checkpointing.{CheckpointFileManager,
CheckpointVersionManager, CommitLog, CommitLogType, CommitMetadataV3,
OffsetLogType, OffsetSeqBase, OffsetSeqLog, OffsetSeqMetadata,
OffsetSeqMetadataBase, OffsetSeqMetadataV2, SinkMetadataInfo}
import
org.apache.spark.sql.execution.streaming.operators.stateful.{StatefulOperatorStateInfo,
StatefulOpStateStoreCheckpointInfo, StateStoreWriter}
import
org.apache.spark.sql.execution.streaming.runtime.StreamingCheckpointConstants.{DIR_NAME_COMMITS,
DIR_NAME_OFFSETS, DIR_NAME_STATE}
import org.apache.spark.sql.execution.streaming.sources.{ForeachBatchSink,
WriteToMicroBatchDataSource, WriteToMicroBatchDataSourceV1}
@@ -55,7 +56,7 @@ import
org.apache.spark.sql.execution.streaming.utils.StreamingUtils
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.internal.connector.PartitionOffsetWithIndex
import org.apache.spark.sql.streaming.Trigger
-import org.apache.spark.util.{Clock, ErrorNotifier, Utils}
+import org.apache.spark.util.{Clock, ErrorNotifier, SerializableConfiguration,
Utils}
class MicroBatchExecution(
sparkSession: SparkSession,
@@ -634,22 +635,6 @@ class MicroBatchExecution(
private def disableAQESupportInStatelessIfUnappropriated(
sparkSessionToRunBatches: SparkSession): Unit = {
- def containsStatefulOperator(p: LogicalPlan): Boolean = {
- p.exists {
- case node: Aggregate if node.isStreaming => true
- case node: Deduplicate if node.isStreaming => true
- case node: DeduplicateWithinWatermark if node.isStreaming => true
- case node: Distinct if node.isStreaming => true
- case node: Join if node.left.isStreaming && node.right.isStreaming =>
true
- case node: FlatMapGroupsWithState if node.isStreaming => true
- case node: FlatMapGroupsInPandasWithState if node.isStreaming => true
- case node: TransformWithState if node.isStreaming => true
- case node: TransformWithStateInPySpark if node.isStreaming => true
- case node: GlobalLimit if node.isStreaming => true
- case _ => false
- }
- }
-
if (trigger.isInstanceOf[RealTimeTrigger]) {
logWarning(log"Disabling AQE since AQE is not supported for Real-time
Mode.")
sparkSessionToRunBatches.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key,
"false")
@@ -660,6 +645,22 @@ class MicroBatchExecution(
}
}
+ private def containsStatefulOperator(p: LogicalPlan): Boolean = {
+ p.exists {
+ case node: Aggregate if node.isStreaming => true
+ case node: Deduplicate if node.isStreaming => true
+ case node: DeduplicateWithinWatermark if node.isStreaming => true
+ case node: Distinct if node.isStreaming => true
+ case node: Join if node.left.isStreaming && node.right.isStreaming =>
true
+ case node: FlatMapGroupsWithState if node.isStreaming => true
+ case node: FlatMapGroupsInPandasWithState if node.isStreaming => true
+ case node: TransformWithState if node.isStreaming => true
+ case node: TransformWithStateInPySpark if node.isStreaming => true
+ case node: GlobalLimit if node.isStreaming => true
+ case _ => false
+ }
+ }
+
/**
* Repeatedly attempts to run batches as data arrives.
*/
@@ -822,20 +823,26 @@ class MicroBatchExecution(
secondLatestOffsets.foreach { offset =>
execCtx.startOffsets = offset.toStreamProgress(sources, sourceIdMap)
}
-
+ val latestCommittedBatch = commitLog.getLatest()
+ val committedBatchId = latestCommittedBatch.map(_._1).getOrElse(-1L)
// update offset metadata
nextOffsets.metadataOpt.foreach { metadata =>
- OffsetSeqMetadata.setSessionConf(metadata,
sparkSessionToRunBatches.sessionState.conf)
+ 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)
+
watermarkTracker.setWatermark(metadataWithRecoveredPartitions.batchWatermarkMs)
}
/* identify the current batch id: if commit log indicates we
successfully processed the
* latest batch id in the offset log, then we can safely move to the
next batch
* i.e., committedBatchId + 1 */
- commitLog.getLatest() match {
+ latestCommittedBatch match {
case Some((latestCommittedBatchId, commitMetadata)) =>
commitMetadata.stateUniqueIds.foreach {
stateUniqueIds => currentStateStoreCkptId ++= stateUniqueIds
@@ -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 and state metadata was
subsequently " +
+ "corrupted due to a bug. See SPARK-59919 for more details.
Delete the " +
+ "checkpoint and restart the query to recover."
+ val numPartitionsOpt = try {
+ val stateMetadataReader = new StateMetadataPartitionReader(
+ stateCheckpointLocation.toString,
+ new SerializableConfiguration(
+ sparkSessionToRunBatches.sessionState.newHadoopConf()),
+ stateMetadataBatchId)
+ stateMetadataReader.stateStoreNumPartitions
+ } catch {
+ case NonFatal(e) => throw new
SparkException(failedRecoveryMessage, e)
+ }
+ val numPartitions = numPartitionsOpt.getOrElse {
+ throw new SparkException(failedRecoveryMessage)
+ }
+ logWarning(log"Recovered state-store partition count " +
+ log"${MDC(CONFIG, SQLConf.SHUFFLE_PARTITIONS.key)}=" +
+ log"${MDC(NUM_PARTITIONS, numPartitions)} from checkpoint state
metadata")
+ OffsetSeqMetadata.withStatefulShufflePartitions(v2, numPartitions)
+ case _ =>
+ metadata
+ }
+ }
+ }
+
/**
* Verify that the checkpoint directory is in a good state to start a new
* streaming query. This checks that the offsets, state, commits directories
are
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/OffsetSeqLogSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/OffsetSeqLogSuite.scala
index ddc789dac8b6..b47c0dc5b372 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/OffsetSeqLogSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/OffsetSeqLogSuite.scala
@@ -19,16 +19,21 @@ package org.apache.spark.sql.execution.streaming
import java.io.File
+import org.apache.hadoop.fs.Path
import org.scalatest.Tag
-import org.apache.spark.sql.AnalysisException
+import org.apache.spark.sql.{AnalysisException, DataFrame}
import org.apache.spark.sql.catalyst.util.stringToFile
-import org.apache.spark.sql.execution.streaming.checkpointing.{OffsetMap,
OffsetSeq, OffsetSeqBase, OffsetSeqLog, OffsetSeqMetadata}
-import org.apache.spark.sql.execution.streaming.runtime.{LongOffset,
MemoryStream, SerializedOffset}
+import
org.apache.spark.sql.execution.datasources.v2.state.metadata.StateMetadataPartitionReader
+import org.apache.spark.sql.execution.streaming.checkpointing.{
+ CommitLog, OffsetMap, OffsetSeq, OffsetSeqBase, OffsetSeqLog,
OffsetSeqMetadata,
+ OffsetSeqMetadataV2}
+import org.apache.spark.sql.execution.streaming.runtime.{
+ LongOffset, MemoryStream, SerializedOffset, StreamingQueryWrapper}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.streaming.StreamingQueryException
import org.apache.spark.sql.test.SharedSparkSession
-import org.apache.spark.util.Utils
+import org.apache.spark.util.{SerializableConfiguration, Utils}
class OffsetSeqLogSuite extends SharedSparkSession {
import testImplicits._
@@ -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") {
+ withSQLConf(
+ SQLConf.SHUFFLE_PARTITIONS.key -> "10",
+ SQLConf.STATEFUL_SHUFFLE_PARTITIONS_INTERNAL.key -> "3") {
+ val v1 = OffsetSeqMetadata(0, 0, spark.conf)
+ val v2 = OffsetSeqMetadataV2(0, 0, spark.conf)
+
+ assert(v1.conf.get(SQLConf.SHUFFLE_PARTITIONS.key).contains("3"))
+ assert(v2.conf.get(SQLConf.SHUFFLE_PARTITIONS.key).contains("3"))
+ }
+ }
+
/** test string offset type */
case class StringOffset(override val json: String) extends Offset
@@ -292,6 +309,285 @@ 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") {
+ withSQLConf(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key -> "1") {
+ 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
+ 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))
+ // The failure is expected because physical state exists for only
five partitions.
+ intercept[StreamingQueryException] {
+ query2.processAllAvailable()
+ }
+ // Planning 10 stores proves the latest offset value won over state
metadata.
+ assert(query2.streamingQuery.lastExecution.numStateStores === 10)
+ } finally {
+ query2.stop()
+ }
+ }
+ }
+ }
+
+ test("SPARK-59919: VERSION_2 recovers state-store partitions when the commit
log is missing") {
+ withSQLConf(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key -> "1") {
+ withTempDir { checkpointDir =>
+ val (_, aggregated) = createStatefulCheckpoint(
+ checkpointDir, "offsetlog_v2_stateful_missing_commit_log_test")
+
+ val offsetLog = new OffsetSeqLog(spark,
s"${checkpointDir.getAbsolutePath}/offsets")
+ updateShufflePartitionsInOffset(
+ offsetLog, checkpointDir, offsetLog.getLatestBatchId().get, None)
+ val hadoopConf = spark.sessionState.newHadoopConf()
+
+ val commitsPath = new Path(checkpointDir.getAbsolutePath, "commits")
+ val commitsFileSystem = commitsPath.getFileSystem(hadoopConf)
+ assert(commitsFileSystem.exists(commitsPath))
+ assert(commitsFileSystem.delete(commitsPath, true))
+ assert(!commitsFileSystem.exists(commitsPath))
+
+ val query2 = startStatefulQuery(
+ aggregated, checkpointDir, shufflePartitions = 10,
+ queryName = "offsetlog_v2_stateful_missing_commit_log_test")
+ try {
+ query2.processAllAvailable()
+ assert(query2.streamingQuery.lastExecution.numStateStores === 5)
+ } finally {
+ query2.stop()
+ }
+ }
+ }
+ }
+
+ test("SPARK-59919: VERSION_2 recovers state-store partitions from metadata
when offset " +
+ "shuffle partitions are missing") {
+ withSQLConf(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key -> "1") {
+ withTempDir { checkpointDir =>
+ val (inputData, aggregated) = createStatefulCheckpoint(
+ checkpointDir, "offsetlog_v2_stateful_missing_partition_test")
+
+ val offsetLog = new OffsetSeqLog(spark,
s"${checkpointDir.getAbsolutePath}/offsets")
+ val latestBatchId = offsetLog.getLatestBatchId().get
+ updateShufflePartitionsInOffset(offsetLog, checkpointDir,
latestBatchId, None)
+ assert(offsetLog.get(latestBatchId).get.metadataOpt.get.conf
+ .get(SQLConf.SHUFFLE_PARTITIONS.key).isEmpty)
+
+ val commitLog = new CommitLog(spark,
s"${checkpointDir.getAbsolutePath}/commits")
+ assert(commitLog.getLatestBatchId().contains(latestBatchId))
+ val stateMetadataReader = new StateMetadataPartitionReader(
+ checkpointDir.getAbsolutePath,
+ new SerializableConfiguration(spark.sessionState.newHadoopConf()),
+ latestBatchId)
+ assert(stateMetadataReader.stateStoreNumPartitions.contains(5))
+
+ // Restart with 10 partitions in session config. Recovery should
restore
+ // the value 5 recorded in state metadata.
+ val query2 = startStatefulQuery(
+ aggregated, checkpointDir, shufflePartitions = 10,
+ queryName = "offsetlog_v2_stateful_missing_partition_test")
+ try {
+ inputData.addData((4, 0))
+ query2.processAllAvailable()
+ assert(query2.streamingQuery.lastExecution.numStateStores === 5)
+ } finally {
+ query2.stop()
+ }
+ }
+ }
+ }
+
+ test("SPARK-59919: VERSION_2 fails when state metadata is corrupt during
partition recovery") {
+ withSQLConf(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key -> "1") {
+ withTempDir { checkpointDir =>
+ val (_, aggregated) = createStatefulCheckpoint(
+ checkpointDir, "offsetlog_v2_stateful_corrupt_metadata_test")
+
+ val offsetLog = new OffsetSeqLog(spark,
s"${checkpointDir.getAbsolutePath}/offsets")
+ updateShufflePartitionsInOffset(
+ offsetLog, checkpointDir, offsetLog.getLatestBatchId().get, None)
+
+ val stateMetadataFile = new File(
+ checkpointDir,
+ "state/0/_metadata/metadata")
+ assert(stateMetadataFile.isFile, s"Missing state metadata file:
$stateMetadataFile")
+ stringToFile(stateMetadataFile, "v1\ncorrupt state metadata")
+
+ assertStatefulPartitionRecoveryFails(
+ aggregated, checkpointDir,
"offsetlog_v2_stateful_corrupt_metadata_test")
+ assert(offsetLog.getLatest().get._2.metadataOpt.get.conf
+ .get(SQLConf.SHUFFLE_PARTITIONS.key).isEmpty)
+ }
+ }
+ }
+
+ test("SPARK-59919: VERSION_2 fails when state metadata is missing during
partition recovery") {
+ withSQLConf(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key -> "1") {
+ withTempDir { checkpointDir =>
+ val (_, aggregated) = createStatefulCheckpoint(
+ checkpointDir, "offsetlog_v2_stateful_missing_metadata_test")
+
+ val offsetLog = new OffsetSeqLog(spark,
s"${checkpointDir.getAbsolutePath}/offsets")
+ updateShufflePartitionsInOffset(
+ offsetLog, checkpointDir, offsetLog.getLatestBatchId().get, None)
+
+ val stateMetadataFile = new File(checkpointDir,
"state/0/_metadata/metadata")
+ assert(stateMetadataFile.isFile, s"Missing state metadata file:
$stateMetadataFile")
+ assert(
+ stateMetadataFile.delete(),
+ s"Failed to delete state metadata file: $stateMetadataFile")
+
+ assertStatefulPartitionRecoveryFails(
+ aggregated, checkpointDir,
"offsetlog_v2_stateful_missing_metadata_test")
+ }
+ }
+ }
+
+ test("SPARK-59919: VERSION_2 uses session shuffle partitions for a stateless
query") {
+ withTempDir { checkpointDir =>
+ withTempDir { outputDir =>
+ val inputData = MemoryStream[Int]
+ val input = inputData.toDF()
+ val queryName = "offsetlog_v2_stateless_missing_partition_test"
+ def startQuery(shufflePartitions: Int): StreamingQueryWrapper = {
+ withSQLConf(
+ SQLConf.STREAMING_OFFSET_LOG_FORMAT_VERSION.key -> "2",
+ SQLConf.SHUFFLE_PARTITIONS.key -> shufflePartitions.toString) {
+ input.writeStream
+ .format("parquet")
+ .queryName(queryName)
+ .option("path", outputDir.getAbsolutePath)
+ .option("checkpointLocation", checkpointDir.getAbsolutePath)
+ .start()
+ .asInstanceOf[StreamingQueryWrapper]
+ }
+ }
+
+ val query1 = startQuery(shufflePartitions = 5)
+ try {
+ inputData.addData(1, 2)
+ query1.processAllAvailable()
+ } finally {
+ query1.stop()
+ }
+
+ val offsetLog = new OffsetSeqLog(spark,
s"${checkpointDir.getAbsolutePath}/offsets")
+ updateShufflePartitionsInOffset(
+ offsetLog, checkpointDir, offsetLog.getLatestBatchId().get, None)
+
+ val query2 = startQuery(shufflePartitions = 10)
+ try {
+ inputData.addData(3, 4)
+ query2.processAllAvailable()
+ assert(offsetLog.getLatest().get._2.metadataOpt.get.conf
+ .get(SQLConf.SHUFFLE_PARTITIONS.key).contains("10"))
+ } finally {
+ query2.stop()
+ }
+ }
+ }
+ }
+
test("STREAMING_OFFSET_LOG_FORMAT_VERSION config - default VERSION_1") {
withTempDir { checkpointDir =>
val inputData = MemoryStream[Int]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
