You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by do...@apache.org on 2026/10/23 12:39:27 UTC
(spark) branch branch-4.x updated: [SPARK-60497][CORE] Reject a map-stage job that depends on a pipelined shuffle
This is an automated email from the ASF dual-hosted git repository.
dongjoon-hyun pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 36e108717195 [SPARK-60497][CORE] Reject a map-stage job that depends on a pipelined shuffle
36e108717195 is described below
commit 36e108717195352e4f801f891f2711d434527864
Author: Dongjoon Hyun <do...@apache.org>
AuthorDate: Fri Oct 23 05:39:10 2026 -0700
[SPARK-60497][CORE] Reject a map-stage job that depends on a pipelined shuffle
### What changes were proposed in this pull request?
This PR makes `DAGScheduler.handleMapStageSubmitted` reject a map-stage job whose shuffle dependency has a `PipelinedShuffleDependency` anywhere in its RDD lineage. The job fails with `PIPELINED_SHUFFLE_UNSUPPORTED` before any stage is created, in the same way that `handleJobSubmitted` rejects the same graph for a result job.
```scala
if (rddGraphHasPipelinedDependency(dependency.rdd)) {
val e = pipelinedUnsupportedError("a map-stage job that depends on a pipelined shuffle")
logWarning(...)
listener.jobFailed(e)
return
}
```
This adds one RDD graph walk per map-stage job, similar to the walk that `handleJobSubmitted` already does for every result job.
### Why are the changes needed?
`PipelinedShuffleDependency` is added newly at 4.3.0.
- #57286
`handleMapStageSubmitted` rejected only a dependency that is itself a `PipelinedShuffleDependency`. A regular `ShuffleDependency` whose map side reads a pipelined shuffle, directly or further down its lineage, was accepted. Such a job skipped every pipelined job check of `handleJobSubmitted`:
- the shuffle shape check, which rejects a pipelined shuffle below a regular shuffle
- gang admission (`rejectUnadmittablePipelinedGroup`)
- `checkPipelinedJobSupported` and the speculation check
- the run epoch (`SPARK_PIPELINED_RUN_EPOCH` and `PipelinedShuffleManager.startRun`)
- `ActiveJob.hasPipelinedDependency`
However, `submitStage` still co-scheduled the pipelined producer with the stage reading it, as ordinary task sets (`TaskSet.isPipelined = false`).
With the in-process channel transport (`PipelinedChannelShuffleManager`), this causes a hang or wrong results.
- **Hang.** The tasks have no run epoch, so every channel access throws `TaskKilledException("Pipelined shuffle run has ended")`. A kill does not count as a task failure in a non-pipelined task set, so the tasks are relaunched forever. In a `local[4]` test, 73,470 task attempts were killed in 10 seconds, each with a WARN log line, and the job never finished.
- **Correctness.** If run epoch 0 is active, because job 0 of the `SparkContext` is a pipelined job that is still running, the channel accesses succeed. A producer attempt that is retried, or relaunched after a kill, then delivers its rows again. The channel reader counts end-of-stream markers without tracking writer attempts, so the map-stage job can succeed with duplicated rows, or stop early and drop the rows of other producers. In local tests, the reader saw 5,120 rows instead of [...]
With the RPC streaming transport, the job runs without gang admission or group-atomic failure handling, and speculation is not rejected.
`SparkContext.submitMapStage` is `private[spark]`, and its only in-tree caller is AQE (`ShuffleExchangeExec`). `AQEEnablePipelinedShuffle` keeps each plan all-pipelined or all-regular, so SQL does not create this shape within one plan. It does when an RDD that contains a pipelined shuffle is wrapped in a new DataFrame and shuffled under AQE. For example, in local mode with `spark.shuffle.manager.incremental=org.apache.spark.shuffle.local.pipelined.PipelinedChannelShuffleManager` and ` [...]
```scala
val inner = spark.range(0, 1000, 1, 2).withColumn("k", $"id" % 7).groupBy($"k").count()
val rows = inner.queryExecution.toRdd.map(r => Row(r.getLong(0), r.getLong(1)))
spark.createDataFrame(rows, inner.schema).groupBy($"count").count().collect()
```
A user RDD that creates a `PipelinedShuffleDependency` (a `DeveloperApi`) reaches it the same way.
### Does this PR introduce _any_ user-facing change?
Yes, within the unreleased branches only. A map-stage job that depends on a pipelined shuffle now fails up front instead of hanging or completing with wrong results. For example, the query above fails with:
```
[PIPELINED_SHUFFLE_UNSUPPORTED] A pipelined shuffle stage group cannot be scheduled because it uses an unsupported feature: a map-stage job that depends on a pipelined shuffle. ...
```
No released Spark version has `PipelinedShuffleDependency`. `branch-4.3` has it with the RPC streaming transport, and `branch-4.x` and `master` also have the in-process channel transport.
### How was this patch tested?
Pass the CIs with the newly added test cases.
- `DAGSchedulerSuite`: the job is rejected when its map stage reads a pipelined shuffle directly and when the pipelined shuffle sits below another regular shuffle, and no task set or partial state is left.
- `PipelinedChannelShuffleSuite`: the job is rejected before any task starts, and the same pipelined shuffle still runs as a result job.
- `AQEPipelinedShuffleSuite`: the query above fails fast instead of hanging, and the inner pipelined query still runs.
Without the fix, the first test fails because the job is accepted, and the other two time out. The duplicated and dropped rows above were reproduced with temporary tests that keep job 0 running, which are not included.
```
$ build/sbt 'core/testOnly *DAGSchedulerSuite *PipelinedChannelShuffleSuite -- -z SPARK-60497'
$ build/sbt 'sql/testOnly *AQEPipelinedShuffleSuite -- -z SPARK-60497'
```
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 5.5
Closes #59738 from dongjoon-hyun/SPARK-60497.
Authored-by: Dongjoon Hyun <do...@apache.org>
Signed-off-by: Dongjoon Hyun <do...@apache.org>
(cherry picked from commit 84dce6db0c8e2dbc86df49a32eb09e5d340c7b61)
Signed-off-by: Dongjoon Hyun <do...@apache.org>
---
.../org/apache/spark/scheduler/DAGScheduler.scala | 12 +++++++
.../apache/spark/scheduler/DAGSchedulerSuite.scala | 29 +++++++++++++++++
.../pipelined/PipelinedChannelShuffleSuite.scala | 27 +++++++++++++++
.../adaptive/AQEPipelinedShuffleSuite.scala | 38 ++++++++++++++++++++++
4 files changed, 106 insertions(+)
diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
index c021d26d2802..45a9948a069b 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
@@ -2825,6 +2825,18 @@ private[spark] class DAGScheduler(
"map output to produce statistics from. This is not supported."))
return
}
+ // A map-stage job also runs its missing ancestor stages, and submitStage co-schedules a
+ // pipelined producer with the stage reading it, but none of handleJobSubmitted's pipelined
+ // setup applies here (shape check, gang admission, idiom checks, run epoch, group-atomic
+ // failure). A pipelined shuffle anywhere below `dependency` is the pipelined-below-regular
+ // shape that handleJobSubmitted rejects, so reject it here too, before any stage is created.
+ if (rddGraphHasPipelinedDependency(dependency.rdd)) {
+ val e = pipelinedUnsupportedError("a map-stage job that depends on a pipelined shuffle")
+ logWarning(log"Rejecting map-stage job ${MDC(JOB_ID, jobId)}: unsupported " +
+ log"pipelined-shuffle idiom", e)
+ listener.jobFailed(e)
+ return
+ }
// Submitting this map stage might still require the creation of some parent stages, so make
// sure that happens.
var finalStage: ShuffleMapStage = null
diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
index f09b7d160ea7..23fc6d602ca5 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
@@ -7441,6 +7441,35 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti
assertDataStructuresEmpty()
}
+ test("SPARK-60497: a map-stage job that depends on a pipelined shuffle is rejected") {
+ // A map-stage job runs its missing ancestor stages, so a pipelined producer below its
+ // dependency would be co-scheduled without any of the pipelined-group setup of a result job.
+ // As in the result-job path, a pipelined shuffle below a regular one is rejected up front,
+ // both when the map stage reads it directly and when it sits below another regular shuffle.
+ // No stage is submitted and no partial state remains.
+ val producerRdd = new MyRDD(sc, 2, Nil)
+ val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2))
+ val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)
+ // consumer --[regular, the map-stage job's dependency]-->
+ val directDep = new ShuffleDependency(consumerRdd, new HashPartitioner(2))
+ // consumer --[regular]--> mid --[regular, the map-stage job's dependency]-->
+ val midRdd = new MyRDD(sc, 2, List(new ShuffleDependency(consumerRdd, new HashPartitioner(2))),
+ tracker = mapOutputTracker)
+ val deepDep = new ShuffleDependency(midRdd, new HashPartitioner(2))
+ for (dep <- Seq(directDep, deepDep)) {
+ val failure = new java.util.concurrent.atomic.AtomicReference[Exception]()
+ val failListener = new JobListener {
+ override def taskSucceeded(index: Int, result: Any): Unit = {}
+ override def jobFailed(exception: Exception): Unit = failure.set(exception)
+ }
+ submitMapStage(dep, listener = failListener)
+ assertPipelinedUnsupported(
+ failure.get(), "a map-stage job that depends on a pipelined shuffle")
+ }
+ assert(taskSets.isEmpty)
+ assertDataStructuresEmpty()
+ }
+
test("regular shuffle map-stage job with speculation enabled is NOT rejected (inertness)") {
// The map-stage speculation fail-fast must be inert for a regular ShuffleDependency.
sc.conf.set(config.SPECULATION_ENABLED, true)
diff --git a/core/src/test/scala/org/apache/spark/shuffle/local/pipelined/PipelinedChannelShuffleSuite.scala b/core/src/test/scala/org/apache/spark/shuffle/local/pipelined/PipelinedChannelShuffleSuite.scala
index a8ff2eee7c7c..d57346879201 100644
--- a/core/src/test/scala/org/apache/spark/shuffle/local/pipelined/PipelinedChannelShuffleSuite.scala
+++ b/core/src/test/scala/org/apache/spark/shuffle/local/pipelined/PipelinedChannelShuffleSuite.scala
@@ -599,6 +599,33 @@ class PipelinedChannelShuffleSuite extends SparkFunSuite {
}
}
+ test("SPARK-60497: a map-stage job that depends on a channel shuffle is rejected up front") {
+ // A map-stage job gets no run epoch, so its channel tasks would be killed on their first
+ // rendezvous access and relaunched forever. The scheduler rejects such a job before any task
+ // runs; the same pipelined shuffle still runs as a result job afterwards.
+ withPipelinedSparkContext(cores = 4) { sc =>
+ val events = new JobEvents
+ sc.addSparkListener(events)
+ val keyed: RDD[(Int, Int)] = sc.parallelize(0 until 100, 2).map(v => (v, v))
+ val consumer = new PipelinedShuffledRDD[Int, Int, Int](keyed, new HashPartitioner(2))
+ val regularDep =
+ new org.apache.spark.ShuffleDependency[Int, Int, Int](consumer, new HashPartitioner(2))
+ val future = sc.submitMapStage(regularDep)
+ org.apache.spark.util.ThreadUtils.awaitReady(
+ future, scala.concurrent.duration.Duration(1, TimeUnit.MINUTES))
+ checkError(
+ exception = future.value.get.failed.get.asInstanceOf[SparkException],
+ condition = "PIPELINED_SHUFFLE_UNSUPPORTED",
+ sqlState = Some("0A000"),
+ parameters = Map("reason" -> "a map-stage job that depends on a pipelined shuffle"))
+ sc.listenerBus.waitUntilEmpty()
+ assert(events.taskStarts.isEmpty)
+ assert(ChannelShuffleRendezvous.numQueuesForTesting === 0)
+
+ assert(consumer.map(_._2).collect().sorted === (0 until 100).toArray)
+ }
+ }
+
test("unregister preserves an active run until its owner finishes") {
ChannelShuffleRendezvous.clear()
try {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AQEPipelinedShuffleSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AQEPipelinedShuffleSuite.scala
index 417c6cdeb4cd..c163231dda02 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AQEPipelinedShuffleSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AQEPipelinedShuffleSuite.scala
@@ -165,4 +165,42 @@ class AQEPipelinedShuffleSuite extends SparkFunSuite
pool.shutdownNow()
}
}
+
+ test("SPARK-60497: an AQE shuffle over a pipelined toRdd is rejected instead of hanging") {
+ withAqePipelinedSession { spark =>
+ import spark.implicits._
+ val inner = spark.range(0, 1000, 1, 2).withColumn("k", ($"id" % 7)).groupBy($"k").count()
+ // QueryExecution.toRdd runs the inner plan as planned: its only exchange is pipelined.
+ val rows = inner.queryExecution.toRdd
+ .map(r => org.apache.spark.sql.Row(r.getLong(0), r.getLong(1)))
+ assert(pipelinedExchanges(inner.queryExecution.executedPlan).nonEmpty)
+ // The outer plan scans an RDD, so its exchange stays regular, and AQE materializes it with a
+ // map-stage job that depends on the inner pipelined shuffle. The scheduler rejects that job
+ // up front. Collect under a deadline so that a regression (the job's channel tasks killed
+ // and relaunched forever) fails here instead of hanging the suite.
+ val outer = spark.createDataFrame(rows, inner.schema).groupBy($"count").count()
+ val pool = Executors.newSingleThreadExecutor()
+ try {
+ val fut = pool.submit(new java.util.concurrent.Callable[Int] {
+ override def call(): Int = outer.collect().length
+ })
+ intercept[java.util.concurrent.ExecutionException] {
+ fut.get(90, TimeUnit.SECONDS)
+ }.getCause match {
+ case e: org.apache.spark.SparkException =>
+ checkError(
+ exception = e,
+ condition = "PIPELINED_SHUFFLE_UNSUPPORTED",
+ sqlState = Some("0A000"),
+ parameters = Map("reason" -> "a map-stage job that depends on a pipelined shuffle"))
+ case other => fail(s"expected a SparkException, got $other", other)
+ }
+ } finally {
+ spark.sparkContext.cancelAllJobs()
+ pool.shutdownNow()
+ }
+ // The inner pipelined query still runs.
+ assert(inner.collect().map(_.getLong(1)).sum === 1000L)
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org