You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ka...@apache.org on 2026/10/12 05:15:29 UTC
(spark) branch master updated: [SPARK-59662][SS] Check for isRepartitionBatch can crash if metadata keys are missing
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 f309aa2a16e1 [SPARK-59662][SS] Check for isRepartitionBatch can crash if metadata keys are missing
f309aa2a16e1 is described below
commit f309aa2a16e111a9a5e3dcdf2751967ce0d36b20
Author: ivasilyev1 <25...@users.noreply.github.com>
AuthorDate: Mon Oct 12 14:15:13 2026 +0900
[SPARK-59662][SS] Check for isRepartitionBatch can crash if metadata keys are missing
### What changes were proposed in this pull request?
This PR changes the repartition checks in `OfflineStateRepartitionUtils.isRepartitionBatch` to be guarded and ensure that both`shufflePartitions` and `previousShufflePartitions` exist before accessing
### Why are the changes needed?
These changes are necessary because in the current implementation `OfflineStateRepartitionUtils.isRepartitionBatch` assumes that `spark.sql.shuffle.partitions` will always be present in the offset metadata and accesses the keys directly with `.get`.
If the offset log is corrupt or otherwise malformed as described in https://issues.apache.org/jira/browse/SPARK-59919 the existing check for `isRepartitionBatch` will cause a hard crash:
```java
26/09/19 14:05:17 ERROR MicroBatchExecution: Query [id = 41da1d72-c034-4a36-9ae3-fdc0e1297405, runId = 4eb3a723-257b-45c7-a2e0-72dff79ebf5b] terminated with error
java.util.NoSuchElementException: None.get
at scala.None$.get(Option.scala:627)
at scala.None$.get(Option.scala:626)
at org.apache.spark.sql.execution.streaming.state.OfflineStateRepartitionUtils$.isRepartitionBatch(OfflineStateRepartitionUtils.scala:54)
at org.apache.spark.sql.execution.streaming.runtime.MicroBatchExecution.checkUnfinishedRepartitionBatch(MicroBatchExecution.scala:564)
at org.apache.spark.sql.execution.streaming.runtime.MicroBatchExecution.initializeExecution(MicroBatchExecution.scala:494)
at org.apache.spark.sql.execution.streaming.runtime.MicroBatchExecution.runActivatedStream(MicroBatchExecution.scala:604)
at org.apache.spark.sql.execution.streaming.runtime.StreamExecution.$anonfun$runStream$1(StreamExecution.scala:353)
at scala.runtime.java8.JFunction0$mcV$sp.apply(JFunction0$mcV$sp.scala:18)
at org.apache.spark.sql.SparkSession.withActive(SparkSession.scala:810)
at org.apache.spark.sql.execution.streaming.runtime.StreamExecution.org$apache$spark$sql$execution$streaming$runtime$StreamExecution$$runStream(StreamExecution.scala:313)
at org.apache.spark.sql.execution.streaming.runtime.StreamExecution$$anon$1.run(StreamExecution.scala:236)
Traceback (most recent call last):
```
This behavior can be recreated with the following script to reliably trigger the crash:
```python
import json
import os
import tempfile
import time
from pathlib import Path
from pyspark.sql import SparkSession
def get_spark():
spark = (
SparkSession.builder
.master("local[2]")
.appName("repartition-crash")
.config("spark.sql.streaming.checkUnfinishedRepartitionOnRestart", "true")
.getOrCreate()
)
return spark
def remove_shuffle_partitions_from_checkpoint(checkpoint_dir):
offsets_dir = Path(checkpoint_dir) / "offsets"
shuffle_partitions_key = "spark.sql.shuffle.partitions"
for offset_file in sorted(
(path for path in offsets_dir.iterdir() if path.name.isdigit()),
key=lambda path: int(path.name),
):
lines = offset_file.read_text().splitlines(keepends=True)
metadata = json.loads(lines[1])
conf = metadata.get("conf")
if shuffle_partitions_key not in conf:
continue
del conf[shuffle_partitions_key]
lines[1] = json.dumps(metadata)
with tempfile.NamedTemporaryFile(
mode="w",
dir=offset_file.parent,
prefix=f".{offset_file.name}.",
suffix=".tmp",
delete=False,
) as temporary_file:
temporary_file.write("".join(lines))
temporary_file.flush()
os.fsync(temporary_file.fileno())
temporary_path = Path(temporary_file.name)
temporary_path.replace(offset_file)
def fail_on_second_batch(batch_df, batch_id):
if batch_id == 1:
raise RuntimeError("Failing on second batch")
batch_df.count()
def main():
input_dir = Path("input")
input_dir.mkdir(exist_ok=True)
( input_dir/"file1.txt" ).write_text("1\n")
( input_dir/"file2.txt" ).write_text("2\n")
spark = get_spark()
df = spark.readStream.format("text").option("maxFilesPerTrigger", "1").load("input")
query = (
df.writeStream
.foreachBatch(fail_on_second_batch)
.option("checkpointLocation", "checkpoint")
.trigger(availableNow=True)
.start()
)
try:
query.awaitTermination()
except:
print("#####################")
print("Expected first run failure")
print("Shutting down spark and recreating")
print("Continue in 5 seconds")
print("#####################")
time.sleep(5)
finally:
query.stop()
spark.stop()
remove_shuffle_partitions_from_checkpoint("checkpoint")
spark = get_spark()
df = (spark.readStream.format("text").option("maxFilesPerTrigger", "1").load("input"))
query = (
df.writeStream
.foreachBatch(lambda batch_df, batch_id: batch_df.count())
.option("checkpointLocation", "checkpoint")
.trigger(availableNow=True)
.start()
)
query.awaitTermination()
if __name__ == "__main__":
main()
```
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
I tested the crash condition with the above python script. After introducing these change and compiling a spark distribution, the query no longer crashes and instead finishes successfully
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: GPT 5.6-Luna
Closes #58929 from ivasilyev1/missing-shuffle-partitions.
Authored-by: ivasilyev1 <25...@users.noreply.github.com>
Signed-off-by: Jungtaek Lim <ka...@gmail.com>
---
.../state/OfflineStateRepartitionUtils.scala | 9 ++++--
.../state/OfflineStateRepartitionSuite.scala | 36 +++++++++++++++++++++-
2 files changed, 41 insertions(+), 4 deletions(-)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OfflineStateRepartitionUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OfflineStateRepartitionUtils.scala
index b1a42ef87bc2..3e5238f2b0ca 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OfflineStateRepartitionUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OfflineStateRepartitionUtils.scala
@@ -51,13 +51,16 @@ object OfflineStateRepartitionUtils {
val previousBatch = prevBatchOpt.get
require(batch.metadataOpt.isDefined, s"Batch $batchId metadata not found")
- val shufflePartitions = getShufflePartitions(batch.metadataOpt.get).get
+ val shufflePartitions = getShufflePartitions(batch.metadataOpt.get)
require(previousBatch.metadataOpt.isDefined,
s"Previous batch $prevBatchId metadata not found")
- val previousShufflePartitions = getShufflePartitions(previousBatch.metadataOpt.get).get
+ val previousShufflePartitions = getShufflePartitions(previousBatch.metadataOpt.get)
- previousShufflePartitions != shufflePartitions
+ (previousShufflePartitions, shufflePartitions) match {
+ case (Some(previous), Some(current)) => previous != current
+ case _ => false
+ }
}
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/OfflineStateRepartitionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/OfflineStateRepartitionSuite.scala
index 2d26121cbdcb..0ba4a81bcdb2 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/OfflineStateRepartitionSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/OfflineStateRepartitionSuite.scala
@@ -23,7 +23,9 @@ import org.apache.hadoop.conf.Configuration
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.execution.datasources.v2.state.metadata.StateMetadataPartitionReader
-import org.apache.spark.sql.execution.streaming.checkpointing.{CommitLog, CommitMetadata}
+import org.apache.spark.sql.execution.streaming.checkpointing.{
+ CommitLog, CommitMetadata, OffsetMap, OffsetSeq, OffsetSeqBase, OffsetSeqLog,
+ OffsetSeqMetadata, OffsetSeqMetadataV2}
import org.apache.spark.sql.execution.streaming.operators.stateful.StatefulOperatorStateInfo
import org.apache.spark.sql.execution.streaming.runtime.{MemoryStream, StreamingQueryCheckpointMetadata}
import org.apache.spark.sql.internal.SQLConf
@@ -108,6 +110,38 @@ class OfflineStateRepartitionSuite extends StreamTest
)
}
+ Seq(1, 2).foreach { offsetLogVersion =>
+ test(s"Missing shuffle partitions metadata is not a repartition batch " +
+ s"(offset log v$offsetLogVersion)") {
+ def assertNotRepartitionBatch(
+ previousConf: Map[String, String],
+ currentConf: Map[String, String]): Unit = {
+ withTempDir { dir =>
+ val offsetLog = new OffsetSeqLog(spark, s"${dir.getAbsolutePath}/offsets")
+
+ def makeOffset(conf: Map[String, String]): OffsetSeqBase = {
+ if (offsetLogVersion == 1) {
+ OffsetSeq(Seq.empty, Some(OffsetSeqMetadata(conf = conf)))
+ } else {
+ OffsetMap(Map.empty, OffsetSeqMetadataV2(conf = conf))
+ }
+ }
+
+ assert(offsetLog.add(0, makeOffset(previousConf)))
+ assert(offsetLog.add(1, makeOffset(currentConf)))
+ assert(!isRepartitionBatch(1, offsetLog))
+ }
+ }
+
+ val shufflePartitionsKey = SQLConf.SHUFFLE_PARTITIONS.key
+ val configuredPartitions = Map(shufflePartitionsKey -> "3")
+
+ assertNotRepartitionBatch(Map.empty, Map.empty)
+ assertNotRepartitionBatch(configuredPartitions, Map.empty)
+ assertNotRepartitionBatch(Map.empty, configuredPartitions)
+ }
+ }
+
Seq(1, 2).foreach { ckptVersion =>
def testWithCheckpointId(testName: String)(testFun: => Unit): Unit = {
test(s"$testName (enableCkptId = ${ckptVersion >= 2})") {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org