From 3231f1e77e3275cd0f1670dd25d9164c5584b45f Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Wed, 22 Jul 2026 04:42:42 +0000 Subject: [PATCH 01/24] [SPARK-XXXXX][CORE] Register a pipelined shuffle with the streaming output tracker A pipelined shuffle streams records from a live producer stage to a co-scheduled consumer stage, so the consumer must look up its producer tasks' locations while both stages run. That routing lives in the StreamingShuffleOutputTracker, but createShuffleMapStage only registered the shuffle with the MapOutputTracker (which, for a transient pipelined shuffle, tracks an empty never-materialized output). As a result a consumer task querying for writer locations found the shuffle unregistered and failed. Register a PipelinedShuffleDependency's shuffle with the StreamingShuffleOutputTracker inside the same once-per-shuffleId guard as the MapOutputTracker registration. Inert for a regular shuffle (and when no streaming tracker is configured). Note: this is a latent gap in the co-scheduling layer independent of RTM; it belongs with the pipelined-scheduling change earlier in the stack and should be re-homed there (with a real-transport regression test) when the stack is next reorganized. Co-authored-by: Isaac --- .../org/apache/spark/scheduler/DAGScheduler.scala | 15 +++++++++++++++ 1 file changed, 15 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 a72af257ad73d..aa0be26f68bbc 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -711,6 +711,21 @@ private[spark] class DAGScheduler( log"shuffle ${MDC(SHUFFLE_ID, shuffleDep.shuffleId)}") mapOutputTracker.registerShuffle(shuffleDep.shuffleId, rdd.partitions.length, shuffleDep.partitioner.numPartitions) + // A pipelined shuffle streams records from a live producer to a co-scheduled consumer, so the + // consumer must be able to look up its producer tasks' locations while both stages run. That + // routing lives in the StreamingShuffleOutputTracker (the MapOutputTracker registration above + // only tracks the empty, never-materialized durable output). Register here, inside the same + // once-per-shuffleId guard, so the tracker knows the producer/consumer fan-in before any + // consumer task asks for a writer location. Inert (tracker absent) for a regular shuffle. + if (shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) { + sc.env.streamingShuffleOutputTracker.foreach { tracker => + tracker.asInstanceOf[StreamingShuffleOutputTrackerMaster].registerShuffle( + shuffleDep.shuffleId, + rdd.partitions.length, + shuffleDep.partitioner.numPartitions, + jobId) + } + } } stage } From d5c04b9eed125ce4c0df73269e7889dbef95bf57 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Mon, 27 Jul 2026 19:15:46 +0000 Subject: [PATCH 02/24] [SPARK-XXXXX][CORE] Unregister a pipelined shuffle from the streaming output tracker on cleanup A pipelined shuffle is registered with the driver-only StreamingShuffleOutputTracker in DAGScheduler.createShuffleMapStage, alongside its MapOutputTracker registration, but nothing ever removed it. Over a long-running Real-Time Mode query -- which allocates a fresh shuffle id per micro-batch -- the tracker's shuffleInfos map grows without bound, and a re-registration of a reused id would hit the "registered twice" guard. Unregister it in ContextCleaner.doCleanupShuffle, right beside the existing mapOutputTrackerMaster.unregisterShuffle. The streaming tracker's state is driver-only (the worker tracker caches nothing and just RPCs the master), so a direct driver-side call is the correct counterpart -- unlike the MapOutputTracker, there is no per-executor cache to invalidate, so this does not belong in the RemoveShuffle RPC handler. The call is guarded by a new containsShuffle predicate so regular (non-pipelined) shuffles, which are never registered here, are skipped without emitting a spurious "not registered" warning on every shuffle cleanup. Also correct the StreamingShuffleManager.unregisterShuffle comment, which previously claimed the tracker was unregistered in BlockManagerStorageEndpoint's RemoveShuffle handler -- untrue in OSS (that path routes to this no-op manager method). Tests: a containsShuffle unit test, plus two ContextCleanerSuite integration tests (a pipelined shuffle is unregistered on cleanup; a regular shuffle is left untouched and quiet). Both integration tests were revert-checked -- each fails when its respective production change (the unregister call, or the containsShuffle guard) is removed. Co-authored-by: Isaac --- .../org/apache/spark/ContextCleaner.scala | 13 ++++++ .../spark/StreamingShuffleOutputTracker.scala | 4 ++ .../streaming/StreamingShuffleManager.scala | 4 +- .../apache/spark/ContextCleanerSuite.scala | 43 +++++++++++++++++++ .../StreamingShuffleOutputTrackerSuite.scala | 14 ++++++ 5 files changed, 76 insertions(+), 2 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/ContextCleaner.scala b/core/src/main/scala/org/apache/spark/ContextCleaner.scala index 0b3c22a22cb46..547112982284a 100644 --- a/core/src/main/scala/org/apache/spark/ContextCleaner.scala +++ b/core/src/main/scala/org/apache/spark/ContextCleaner.scala @@ -242,6 +242,17 @@ private[spark] class ContextCleaner( // to find blocks served by the shuffle service on deallocated executors shuffleDriverComponents.removeShuffle(shuffleId, blocking) mapOutputTrackerMaster.unregisterShuffle(shuffleId) + // A pipelined shuffle is also registered with the driver-only + // StreamingShuffleOutputTracker (see DAGScheduler.createShuffleMapStage). Its state lives + // solely on the driver -- the worker tracker caches nothing -- so unregister it directly + // here, alongside the MapOutputTracker cleanup, rather than through the RemoveShuffle RPC. + // Guarded by containsShuffle so regular (non-pipelined) shuffles are skipped without a + // spurious "not registered" warning. + streamingShuffleOutputTrackerMaster.foreach { tracker => + if (tracker.containsShuffle(shuffleId)) { + tracker.unregisterShuffle(shuffleId) + } + } listeners.asScala.foreach(_.shuffleCleaned(shuffleId)) logDebug("Cleaned shuffle " + shuffleId) } else { @@ -308,6 +319,8 @@ private[spark] class ContextCleaner( private def broadcastManager = sc.env.broadcastManager private def mapOutputTrackerMaster = sc.env.mapOutputTracker.asInstanceOf[MapOutputTrackerMaster] + private def streamingShuffleOutputTrackerMaster: Option[StreamingShuffleOutputTrackerMaster] = + sc.env.streamingShuffleOutputTracker.map(_.asInstanceOf[StreamingShuffleOutputTrackerMaster]) } private object ContextCleaner { diff --git a/core/src/main/scala/org/apache/spark/StreamingShuffleOutputTracker.scala b/core/src/main/scala/org/apache/spark/StreamingShuffleOutputTracker.scala index 090be7f691d4b..9d6fc4f1936c8 100644 --- a/core/src/main/scala/org/apache/spark/StreamingShuffleOutputTracker.scala +++ b/core/src/main/scala/org/apache/spark/StreamingShuffleOutputTracker.scala @@ -230,6 +230,10 @@ private[spark] class StreamingShuffleOutputTrackerMaster(conf: SparkConf) } } + // Whether the given shuffle is registered. ContextCleaner uses this on the driver to skip + // regular (non-pipelined) shuffles, which are never registered here, before unregistering. + def containsShuffle(shuffleId: Int): Boolean = shuffleInfos.containsKey(shuffleId) + // for testing purposes private[spark] def getShuffleInfo(shuffleId: Int): Option[StreamingShuffleInfo] = { Option(shuffleInfos.get(shuffleId)) diff --git a/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala b/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala index 03ca257353c59..b83bd33552b90 100644 --- a/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala +++ b/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala @@ -113,8 +113,8 @@ private[spark] class StreamingShuffleManager extends PipelinedShuffleManager wit override def unregisterShuffle(shuffleId: Int): Boolean = { // No manager-side state to release here: the driver's StreamingShuffleOutputTracker is - // unregistered in BlockManagerStorageEndpoint's RemoveShuffle handler, and per-task writer - // and reader resources are released via task completion listeners. + // unregistered in ContextCleaner.doCleanupShuffle (its state is driver-only), and per-task + // writer and reader resources are released via task completion listeners. true } diff --git a/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala b/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala index bb2d7d5c4d8e4..7b00e4440cce8 100644 --- a/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala +++ b/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala @@ -20,6 +20,7 @@ package org.apache.spark import scala.collection.mutable.HashSet import scala.util.Random +import org.apache.logging.log4j.Level import org.scalatest.BeforeAndAfter import org.scalatest.concurrent.Eventually._ import org.scalatest.concurrent.PatienceConfiguration @@ -127,6 +128,48 @@ class ContextCleanerSuite extends ContextCleanerSuiteBase { assert(rdd.collect().toList.equals(collected)) } + test("cleanup shuffle unregisters a pipelined shuffle from the streaming output tracker") { + // A pipelined shuffle is registered in BOTH the MapOutputTracker and the driver-only + // StreamingShuffleOutputTracker (see DAGScheduler.createShuffleMapStage). doCleanupShuffle + // must unregister it from the streaming tracker too, mirroring the MapOutputTracker cleanup; + // otherwise the tracker grows without bound across micro-batches in Real-Time Mode. + val streamingTracker = sc.env.streamingShuffleOutputTracker.get + .asInstanceOf[StreamingShuffleOutputTrackerMaster] + val mapOutputTracker = sc.env.mapOutputTracker.asInstanceOf[MapOutputTrackerMaster] + + // Register a shuffle id in both trackers exactly as the scheduler does for a pipelined shuffle. + val shuffleId = 1000 + mapOutputTracker.registerShuffle(shuffleId, numMaps = 2, numReduces = 2) + streamingTracker.registerShuffle(shuffleId, numMaps = 2, numReduces = 2, jobId = 0) + assert(streamingTracker.containsShuffle(shuffleId)) + + cleaner.doCleanupShuffle(shuffleId, blocking = true) + + // The streaming tracker entry is gone, alongside the MapOutputTracker cleanup. + assert(!streamingTracker.containsShuffle(shuffleId)) + assert(!mapOutputTracker.containsShuffle(shuffleId)) + } + + test("cleanup shuffle leaves a regular shuffle's streaming tracker untouched and quiet") { + // A regular (non-pipelined) shuffle is never registered with the streaming tracker, so + // doCleanupShuffle must skip it via the containsShuffle guard -- without emitting the + // "attempting to unregister a shuffle that hasn't been registered" warning for every shuffle. + val streamingTracker = sc.env.streamingShuffleOutputTracker.get + .asInstanceOf[StreamingShuffleOutputTrackerMaster] + + val (rdd, shuffleDeps) = newRDDWithShuffleDependencies() + rdd.collect() + shuffleDeps.foreach(s => assert(!streamingTracker.containsShuffle(s.shuffleId))) + + val logAppender = new LogAppender("unregister streaming shuffle") + logAppender.setThreshold(Level.WARN) + withLogAppender(logAppender, level = Some(Level.WARN)) { + shuffleDeps.foreach(s => cleaner.doCleanupShuffle(s.shuffleId, blocking = true)) + } + assert(!logAppender.loggingEvents.exists( + _.getMessage.getFormattedMessage.contains("hasn't been registered"))) + } + test("cleanup broadcast") { val broadcast = newBroadcast() val tester = new CleanerTester(sc, broadcastIds = Seq(broadcast.id)) diff --git a/core/src/test/scala/org/apache/spark/StreamingShuffleOutputTrackerSuite.scala b/core/src/test/scala/org/apache/spark/StreamingShuffleOutputTrackerSuite.scala index 3c8d9b2f14e21..a9351630ad8db 100644 --- a/core/src/test/scala/org/apache/spark/StreamingShuffleOutputTrackerSuite.scala +++ b/core/src/test/scala/org/apache/spark/StreamingShuffleOutputTrackerSuite.scala @@ -270,6 +270,20 @@ class StreamingShuffleOutputTrackerSuite assert(tracker.getAllShuffleWriterTaskLocations(0).isEmpty) } + test("StreamingShuffleOutputTrackerMaster - containsShuffle reflects register/unregister") { + val conf = new SparkConf(false) + val tracker = newTrackerMaster(conf) + + // A shuffle that was never registered is absent. + assert(!tracker.containsShuffle(0)) + + tracker.registerShuffle(shuffleId = 0, numMaps = 1, numReduces = 1, jobId = 1) + assert(tracker.containsShuffle(0)) + + tracker.unregisterShuffle(0) + assert(!tracker.containsShuffle(0)) + } + test("StreamingShuffleOutputTrackerMaster - register writer before shuffle fails") { val conf = new SparkConf(false) val tracker = newTrackerMaster(conf) From 309498b67acd6898a4648b2c397329a7628b1c2f Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Tue, 28 Jul 2026 02:10:35 +0000 Subject: [PATCH 03/24] [SPARK-XXXXX][CORE] Fully separate pipelined shuffles from the MapOutputTracker A pipelined shuffle produces no durable, addressable map output: its producer/consumer task-location routing lives entirely in the StreamingShuffleOutputTracker, its availability is tracked on the stage (pipelinedCompletedPartitions), and its reader locates writers through the streaming tracker. Yet createShuffleMapStage still registered an empty ShuffleStatus for it in the MapOutputTracker, purely to keep containsShuffle / unregisterShuffle bookkeeping uniform. That coupled the streaming tracker's lifecycle to MapOutputTracker state: registration and cleanup of the streaming entry were both nested inside MapOutputTracker.containsShuffle guards, so a future divergence between the two trackers could double-register (throwing "registered twice") or leak the streaming entry. Split the two trackers cleanly by dependency type, with no overlap: - createShuffleMapStage: a PipelinedShuffleDependency registers ONLY with the StreamingShuffleOutputTracker (self-guarded on that tracker's own containsShuffle); a regular shuffle registers ONLY with the MapOutputTracker, exactly as before. No empty MapOutputTracker shell is created for a pipelined shuffle. - ContextCleaner.doCleanupShuffle: two independent branches, each keyed on its own tracker's membership -- MapOutputTracker cleanup for a regular shuffle, a direct driver-side StreamingShuffleOutputTracker unregister for a pipelined one (no RemoveShuffle RPC: a streaming shuffle has no block-manager-served files). Neither branch depends on the other tracker's state. This is safe because every MapOutputTracker path a pipelined shuffle could reach is either already type-gated out (output registration, availability, FetchFailed invalidation which routes to a group-atomic abort, push-merge, indeterminate/checksum rollback, map-stage-job stats) or tolerant of the entry's absence (executor/host-loss strips iterate shuffleStatuses and simply skip it). Consumer task locality never reads the tracker for a shuffle dependency (getPreferredLocsInternal recurses only through narrow deps). Tests (ContextCleanerSuite): a pipelined shuffle registered ONLY in the streaming tracker is still cleaned up by doCleanupShuffle even though it is absent from the MapOutputTracker (revert-checked: fails without the streaming cleanup branch); a regular shuffle takes the MapOutputTracker branch and never touches the streaming tracker. Co-authored-by: Isaac --- .../org/apache/spark/ContextCleaner.scala | 23 +++++----- .../apache/spark/scheduler/DAGScheduler.scala | 42 ++++++++++++------- .../spark/scheduler/ShuffleMapStage.scala | 13 +++--- .../apache/spark/ContextCleanerSuite.scala | 26 +++++++----- 4 files changed, 59 insertions(+), 45 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/ContextCleaner.scala b/core/src/main/scala/org/apache/spark/ContextCleaner.scala index 547112982284a..a01a480147d72 100644 --- a/core/src/main/scala/org/apache/spark/ContextCleaner.scala +++ b/core/src/main/scala/org/apache/spark/ContextCleaner.scala @@ -236,25 +236,26 @@ private[spark] class ContextCleaner( /** Perform shuffle cleanup. */ def doCleanupShuffle(shuffleId: Int, blocking: Boolean): Unit = { try { + // A shuffle lives in exactly one tracker, split by dependency type: a regular shuffle in the + // MapOutputTracker, a pipelined shuffle in the driver-only StreamingShuffleOutputTracker (see + // DAGScheduler.createShuffleMapStage). Clean up whichever holds it -- the two branches are + // independent, each keyed on its own tracker's membership, so neither depends on the other. if (mapOutputTrackerMaster.containsShuffle(shuffleId)) { logDebug("Cleaning shuffle " + shuffleId) // Shuffle must be removed before it's unregistered from the output tracker // to find blocks served by the shuffle service on deallocated executors shuffleDriverComponents.removeShuffle(shuffleId, blocking) mapOutputTrackerMaster.unregisterShuffle(shuffleId) - // A pipelined shuffle is also registered with the driver-only - // StreamingShuffleOutputTracker (see DAGScheduler.createShuffleMapStage). Its state lives - // solely on the driver -- the worker tracker caches nothing -- so unregister it directly - // here, alongside the MapOutputTracker cleanup, rather than through the RemoveShuffle RPC. - // Guarded by containsShuffle so regular (non-pipelined) shuffles are skipped without a - // spurious "not registered" warning. - streamingShuffleOutputTrackerMaster.foreach { tracker => - if (tracker.containsShuffle(shuffleId)) { - tracker.unregisterShuffle(shuffleId) - } - } listeners.asScala.foreach(_.shuffleCleaned(shuffleId)) logDebug("Cleaned shuffle " + shuffleId) + } else if (streamingShuffleOutputTrackerMaster.exists(_.containsShuffle(shuffleId))) { + // A pipelined shuffle's state is driver-only (the worker tracker caches nothing), so + // cleanup is a direct driver-side unregister. No RemoveShuffle RPC is needed: a streaming + // shuffle has no durable, block-manager-served files to remove. + logDebug("Cleaning pipelined shuffle " + shuffleId) + streamingShuffleOutputTrackerMaster.foreach(_.unregisterShuffle(shuffleId)) + listeners.asScala.foreach(_.shuffleCleaned(shuffleId)) + logDebug("Cleaned pipelined shuffle " + shuffleId) } else { logDebug("Asked to cleanup non-existent shuffle (maybe it was already removed)") } 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 aa0be26f68bbc..c756d8d3c8263 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -703,7 +703,32 @@ private[spark] class DAGScheduler( shuffleIdToMapStage(shuffleDep.shuffleId) = stage updateJobIdStageIdMaps(jobId, stage) - if (!mapOutputTracker.containsShuffle(shuffleDep.shuffleId)) { + // A pipelined shuffle and a regular shuffle live in DIFFERENT trackers -- the split is by + // dependency type, with no overlap. A pipelined shuffle produces no durable, addressable map + // output; its producer/consumer task-location routing lives entirely in the + // StreamingShuffleOutputTracker, and it is never registered with the MapOutputTracker (nothing + // reads such an entry -- availability is tracked on the stage via pipelinedCompletedPartitions, + // the reader locates writers through the streaming tracker, and every MapOutputTracker path a + // pipelined shuffle could reach is either type-gated out or tolerant of its absence). A regular + // shuffle is registered with the MapOutputTracker exactly as before. + if (shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) { + // Register the producer/consumer fan-in before any consumer task asks for a writer location. + // Self-guarded on the streaming tracker's own state (createShuffleMapStage runs once per + // shuffleId via getOrCreateShuffleMapStage, but guard defensively against a re-entry). + sc.env.streamingShuffleOutputTracker.foreach { tracker => + val streamingTracker = tracker.asInstanceOf[StreamingShuffleOutputTrackerMaster] + if (!streamingTracker.containsShuffle(shuffleDep.shuffleId)) { + logInfo(log"Registering RDD ${MDC(RDD_ID, rdd.id)} " + + log"(${MDC(CREATION_SITE, rdd.getCreationSite)}) as input to pipelined " + + log"shuffle ${MDC(SHUFFLE_ID, shuffleDep.shuffleId)}") + streamingTracker.registerShuffle( + shuffleDep.shuffleId, + rdd.partitions.length, + shuffleDep.partitioner.numPartitions, + jobId) + } + } + } else if (!mapOutputTracker.containsShuffle(shuffleDep.shuffleId)) { // Kind of ugly: need to register RDDs with the cache and map output tracker here // since we can't do it in the RDD constructor because # of partitions is unknown logInfo(log"Registering RDD ${MDC(RDD_ID, rdd.id)} " + @@ -711,21 +736,6 @@ private[spark] class DAGScheduler( log"shuffle ${MDC(SHUFFLE_ID, shuffleDep.shuffleId)}") mapOutputTracker.registerShuffle(shuffleDep.shuffleId, rdd.partitions.length, shuffleDep.partitioner.numPartitions) - // A pipelined shuffle streams records from a live producer to a co-scheduled consumer, so the - // consumer must be able to look up its producer tasks' locations while both stages run. That - // routing lives in the StreamingShuffleOutputTracker (the MapOutputTracker registration above - // only tracks the empty, never-materialized durable output). Register here, inside the same - // once-per-shuffleId guard, so the tracker knows the producer/consumer fan-in before any - // consumer task asks for a writer location. Inert (tracker absent) for a regular shuffle. - if (shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) { - sc.env.streamingShuffleOutputTracker.foreach { tracker => - tracker.asInstanceOf[StreamingShuffleOutputTrackerMaster].registerShuffle( - shuffleDep.shuffleId, - rdd.partitions.length, - shuffleDep.partitioner.numPartitions, - jobId) - } - } } stage } diff --git a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala index 94078d698cf4f..135e5e5a4bf2c 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala @@ -64,13 +64,12 @@ private[spark] class ShuffleMapStage( /** * Availability tracking for a pipelined shuffle. A pipelined shuffle produces no durable, - * addressable map output. `createShuffleMapStage` still calls `MapOutputTracker.registerShuffle` - * for it (it is a `ShuffleDependency`, and an empty shuffle-status entry keeps the tracker's - * `containsShuffle` / `unregisterShuffle` bookkeeping uniform), but no map output is ever - * registered into that entry -- `registerMapOutput` runs only on the non-pipelined completion - * path -- so the entry stays empty. Its map-stage availability therefore cannot be read from the - * tracker; instead we track completed partitions here, monotonically: a partition is added when - * its map task succeeds and is NEVER removed on executor/host loss. + * addressable map output and is NOT registered with the `MapOutputTracker` at all -- + * `createShuffleMapStage` registers it only with the `StreamingShuffleOutputTracker` (the two + * trackers are split by dependency type, with no overlap). Its map-stage availability therefore + * cannot be read from the `MapOutputTracker`; instead we track completed partitions here, + * monotonically: a partition is added when its map task succeeds and is NEVER removed on + * executor/host loss. * * This is the crux of avoiding the streaming-writer resubmit hang: if a pipelined shuffle's * availability were read from the `MapOutputTracker`, losing an executor that held a completed diff --git a/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala b/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala index 7b00e4440cce8..e8f33b7803da9 100644 --- a/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala +++ b/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala @@ -129,31 +129,35 @@ class ContextCleanerSuite extends ContextCleanerSuiteBase { } test("cleanup shuffle unregisters a pipelined shuffle from the streaming output tracker") { - // A pipelined shuffle is registered in BOTH the MapOutputTracker and the driver-only - // StreamingShuffleOutputTracker (see DAGScheduler.createShuffleMapStage). doCleanupShuffle - // must unregister it from the streaming tracker too, mirroring the MapOutputTracker cleanup; - // otherwise the tracker grows without bound across micro-batches in Real-Time Mode. + // A pipelined shuffle lives ONLY in the driver-only StreamingShuffleOutputTracker -- it is + // never registered with the MapOutputTracker (see DAGScheduler.createShuffleMapStage). So the + // MapOutputTracker.containsShuffle guard is false for it, and doCleanupShuffle must still clean + // it up via its own streaming branch; otherwise the tracker grows without bound across + // micro-batches in Real-Time Mode. val streamingTracker = sc.env.streamingShuffleOutputTracker.get .asInstanceOf[StreamingShuffleOutputTrackerMaster] val mapOutputTracker = sc.env.mapOutputTracker.asInstanceOf[MapOutputTrackerMaster] - // Register a shuffle id in both trackers exactly as the scheduler does for a pipelined shuffle. + // Register a shuffle id in the streaming tracker ONLY, exactly as the scheduler now does for a + // pipelined shuffle (nothing is put in the MapOutputTracker). val shuffleId = 1000 - mapOutputTracker.registerShuffle(shuffleId, numMaps = 2, numReduces = 2) streamingTracker.registerShuffle(shuffleId, numMaps = 2, numReduces = 2, jobId = 0) assert(streamingTracker.containsShuffle(shuffleId)) + assert(!mapOutputTracker.containsShuffle(shuffleId), + "a pipelined shuffle must not be registered with the MapOutputTracker") cleaner.doCleanupShuffle(shuffleId, blocking = true) - // The streaming tracker entry is gone, alongside the MapOutputTracker cleanup. + // The streaming tracker entry is gone -- cleanup fired even though the shuffle was never in the + // MapOutputTracker, proving the streaming cleanup branch is independent of MapOutputTracker. assert(!streamingTracker.containsShuffle(shuffleId)) - assert(!mapOutputTracker.containsShuffle(shuffleId)) } test("cleanup shuffle leaves a regular shuffle's streaming tracker untouched and quiet") { - // A regular (non-pipelined) shuffle is never registered with the streaming tracker, so - // doCleanupShuffle must skip it via the containsShuffle guard -- without emitting the - // "attempting to unregister a shuffle that hasn't been registered" warning for every shuffle. + // A regular (non-pipelined) shuffle is registered only with the MapOutputTracker, never the + // streaming tracker. doCleanupShuffle must take the MapOutputTracker branch and never touch the + // streaming tracker -- in particular it must not emit the streaming tracker's "attempting to + // unregister a shuffle that hasn't been registered" warning. val streamingTracker = sc.env.streamingShuffleOutputTracker.get .asInstanceOf[StreamingShuffleOutputTrackerMaster] From 2bd06ffc41eca5103ff0c8220672506e64280295 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Tue, 28 Jul 2026 02:48:45 +0000 Subject: [PATCH 04/24] [SPARK-XXXXX][CORE] Harden the MapOutputTracker decoupling: fail loud (before any state mutation) with no streaming tracker, and abort the group on a base-path pipelined FetchFailed Three hardening fixes to "Fully separate pipelined shuffles from the MapOutputTracker", all found by review of that change: 1. createShuffleMapStage registered a pipelined dependency only inside a streamingShuffleOutputTracker.foreach, with the MapOutputTracker registration in the else-if branch. If the streaming tracker were absent (a non-streaming custom incremental manager), the shuffle was registered in NEITHER tracker -- a consumer would then fail with an opaque NullPointerException in the shuffle writer. Fail loud instead: a PipelinedShuffleDependency requires a StreamingShuffleOutputTracker. Crucially, resolve the tracker UP FRONT -- before createShuffleMapStage mutates stageIdToStage / shuffleIdToMapStage / updateJobIdStageIdMaps -- so the fail-loud throw leaves no partial scheduler state. (handleJobSubmitted's catch does no rollback, so a throw after those mutations would leak a half-created stage; a re-submission would then reuse the stale cached stage and fail with a misleading cross-job-reuse error instead of re-throwing the real misconfiguration.) 2. A pipelined-group FetchFailed is normally intercepted by the group-atomic abort branch, whose guard is keyed on the failed stage's active job carrying hasPipelinedDependency. If a pipelined-shuffle FetchFailed reaches the base (non-group-atomic) path instead, that path is invalid for a pipelined shuffle in two ways: (a) its MapOutputTracker invalidation would throw ShuffleStatusNotFoundException (a pipelined shuffle is no longer registered there); and (b) its resubmit branch would enqueue a lone-stage resubmit of the transient producer, which cannot be re-read and would deadlock the group. Guard the whole base-path tail on mapStage.isPipelined and abort the group instead -- matching the group-atomic branch's outcome. Tests (revert-checked): - PipelinedShuffleRoutingSuite: a job with a pipelined dependency and no streaming tracker fails loud with the IllegalStateException, and a re-submission re-throws the SAME error -- proving no partial stage leaked (revert-checked against the throw-after-mutation placement, which instead leaks and yields a cross-job-reuse error on re-submit). - DAGSchedulerSuite: a FetchFailed for a pipelined shuffle reaching the base path (group-atomic guard forced to miss) aborts the group -- the job fails, the throwing MapOutputTracker calls are never invoked, and the transient producer is not resubmitted (revert-checked against both the no-guard and the invalidation-only-guard variants). Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 240 +++++++++++------- .../PipelinedShuffleRoutingSuite.scala | 40 +++ 2 files changed, 184 insertions(+), 96 deletions(-) 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 c756d8d3c8263..b64a9f5282aca 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -692,6 +692,24 @@ private[spark] class DAGScheduler( checkBarrierStageWithNumSlots(rdd, resourceProfile) checkBarrierStageWithRDDChainPattern(rdd, rdd.getNumPartitions) checkPipelinedProducerSupported(shuffleDep) + // A pipelined shuffle is served solely by the StreamingShuffleOutputTracker, so that tracker + // MUST exist for one -- it is created whenever a streaming-capable shuffle manager is set up + // (see SparkEnv.initializeStreamingShuffleOutputTracker), a prerequisite for producing a + // PipelinedShuffleDependency. Resolve it up front (BEFORE any stage-map mutation below), so a + // fail-loud on the misconfiguration leaves no partial scheduler state -- and so a pipelined + // shuffle is never silently registered in no tracker (a consumer would then find no writer + // locations). The reader enforces the same invariant (see StreamingShuffleReader). None for a + // regular shuffle, which is registered with the MapOutputTracker instead. + val streamingTrackerForPipelined: Option[StreamingShuffleOutputTrackerMaster] = + if (shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) { + Some(sc.env.streamingShuffleOutputTracker + .getOrElse(throw new IllegalStateException( + s"A pipelined shuffle (id ${shuffleDep.shuffleId}) requires a " + + "StreamingShuffleOutputTracker, but none is configured")) + .asInstanceOf[StreamingShuffleOutputTrackerMaster]) + } else { + None + } val numTasks = rdd.partitions.length val parents = getOrCreateParentStages(shuffleDeps, jobId) val id = nextStageId.getAndIncrement() @@ -706,17 +724,22 @@ private[spark] class DAGScheduler( // A pipelined shuffle and a regular shuffle live in DIFFERENT trackers -- the split is by // dependency type, with no overlap. A pipelined shuffle produces no durable, addressable map // output; its producer/consumer task-location routing lives entirely in the - // StreamingShuffleOutputTracker, and it is never registered with the MapOutputTracker (nothing - // reads such an entry -- availability is tracked on the stage via pipelinedCompletedPartitions, - // the reader locates writers through the streaming tracker, and every MapOutputTracker path a - // pipelined shuffle could reach is either type-gated out or tolerant of its absence). A regular - // shuffle is registered with the MapOutputTracker exactly as before. - if (shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) { - // Register the producer/consumer fan-in before any consumer task asks for a writer location. - // Self-guarded on the streaming tracker's own state (createShuffleMapStage runs once per - // shuffleId via getOrCreateShuffleMapStage, but guard defensively against a re-entry). - sc.env.streamingShuffleOutputTracker.foreach { tracker => - val streamingTracker = tracker.asInstanceOf[StreamingShuffleOutputTrackerMaster] + // StreamingShuffleOutputTracker (resolved fail-loud above), and it is never registered with the + // MapOutputTracker (nothing reads such an entry -- availability is tracked on the stage via + // pipelinedCompletedPartitions, and the reader locates writers through the streaming tracker). + // The MapOutputTracker.getStatistics paths, which WOULD throw ShuffleStatusNotFoundException on + // the absent entry, are all unreachable for a pipelined dependency: markMapStageJobsAsFinished + // only calls it when mapStageJobs is non-empty, but handleMapStageSubmitted rejects a pipelined + // dependency before addActiveJob (its sole populator) ever runs, so a pipelined stage's + // mapStageJobs is always empty; and the getStatistics in checkAndScheduleShuffleMergeFinalize + // is on the push-based-shuffle merge path, which a pipelined dependency rejects up front + // (shuffleMergeEnabled in checkPipelinedProducerSupported). A regular shuffle registers with + // the MapOutputTracker as before. + streamingTrackerForPipelined match { + case Some(streamingTracker) => + // Register the producer/consumer fan-in before any consumer task asks for a writer loc. + // Self-guarded on the tracker's own state (createShuffleMapStage runs once per shuffleId + // via getOrCreateShuffleMapStage, but guard defensively against a re-entry). if (!streamingTracker.containsShuffle(shuffleDep.shuffleId)) { logInfo(log"Registering RDD ${MDC(RDD_ID, rdd.id)} " + log"(${MDC(CREATION_SITE, rdd.getCreationSite)}) as input to pipelined " + @@ -727,15 +750,16 @@ private[spark] class DAGScheduler( shuffleDep.partitioner.numPartitions, jobId) } - } - } else if (!mapOutputTracker.containsShuffle(shuffleDep.shuffleId)) { - // Kind of ugly: need to register RDDs with the cache and map output tracker here - // since we can't do it in the RDD constructor because # of partitions is unknown - logInfo(log"Registering RDD ${MDC(RDD_ID, rdd.id)} " + - log"(${MDC(CREATION_SITE, rdd.getCreationSite)}) as input to " + - log"shuffle ${MDC(SHUFFLE_ID, shuffleDep.shuffleId)}") - mapOutputTracker.registerShuffle(shuffleDep.shuffleId, rdd.partitions.length, - shuffleDep.partitioner.numPartitions) + case None if !mapOutputTracker.containsShuffle(shuffleDep.shuffleId) => + // Kind of ugly: need to register RDDs with the cache and map output tracker here + // since we can't do it in the RDD constructor because # of partitions is unknown + logInfo(log"Registering RDD ${MDC(RDD_ID, rdd.id)} " + + log"(${MDC(CREATION_SITE, rdd.getCreationSite)}) as input to " + + log"shuffle ${MDC(SHUFFLE_ID, shuffleDep.shuffleId)}") + mapOutputTracker.registerShuffle(shuffleDep.shuffleId, rdd.partitions.length, + shuffleDep.partitioner.numPartitions) + + case _ => // regular shuffle already registered with the MapOutputTracker; nothing to do } stage } @@ -3071,8 +3095,17 @@ private[spark] class DAGScheduler( // will be re-executed. if (clearShuffle) { logInfo(log"Cleaning up shuffle for stage ${MDC(STAGE, sms)} to ensure re-execution") - mapOutputTracker.unregisterAllMapAndMergeOutput(sms.shuffleDep.shuffleId) - sms.shuffleDep.newShuffleMergeState() + // A pipelined shuffle is not registered with the MapOutputTracker, so unregistering there + // would throw ShuffleStatusNotFoundException. Not reachable today -- an indeterminate + // pipelined producer is rejected up front (checkPipelinedProducerSupported), and a job is + // all-regular or all-pipelined (mixed rejected), so a pipelined stage is never a succeeding + // stage of a regular indeterminate producer that rolls back -- but guard defensively, like + // the pipelined branch on the FetchFailed base path. (A transient pipelined producer cannot + // be rolled back and recomputed anyway; a genuine member failure fails the group, not this.) + if (!sms.isPipelined) { + mapOutputTracker.unregisterAllMapAndMergeOutput(sms.shuffleDep.shuffleId) + sms.shuffleDep.newShuffleMergeState() + } } } @@ -3436,86 +3469,101 @@ private[spark] class DAGScheduler( "longer running") } - if (mapStage.rdd.isBarrier()) { - // Mark all the map as broken in the map stage, to ensure retry all the tasks on - // resubmitted stage attempt. - // TODO: SPARK-35547: Clean all push-based shuffle metadata like merge enabled and - // TODO: finalized as we are clearing all the merge results. - mapOutputTracker.unregisterAllMapAndMergeOutput(shuffleId) - } else if (mapIndex != -1) { - // Mark the map whose fetch failed as broken in the map stage - mapOutputTracker.unregisterMapOutput(shuffleId, mapIndex, bmAddress) - if (pushBasedShuffleEnabled) { - // Possibly unregister the merge result , if the FetchFailed - // mapIndex is part of the merge result of - mapOutputTracker. - unregisterMergeResult(shuffleId, reduceId, bmAddress, Option(mapIndex)) - } + if (mapStage.isPipelined) { + // Defense-in-depth for a pipelined-shuffle FetchFailed that reaches this base path + // rather than the group-atomic branch above (e.g. a job whose hasPipelinedDependency + // flag was not propagated through the group check). The base path is invalid for a + // pipelined shuffle in two ways: (a) the MapOutputTracker invalidation below would + // throw ShuffleStatusNotFoundException (a pipelined shuffle is never registered there, + // see createShuffleMapStage); (b) the resubmit branch would enqueue a lone-stage + // resubmit of the transient producer, which -- as the group-atomic branch's comment + // explains -- is never valid and would deadlock the group. So abort the group here + // instead, matching the group-atomic branch's outcome, and skip both. + abortStage(failedStage, + s"A pipelined group member failed with a fetch failure: $failureMessage", None) } else { - // Unregister the merge result of if there is a FetchFailed event - // and is not a MetaDataFetchException which is signified by bmAddress being null - if (bmAddress != null && - bmAddress.executorId.equals(BlockManagerId.SHUFFLE_MERGER_IDENTIFIER)) { - assert(pushBasedShuffleEnabled, "Push based shuffle expected to " + - "be enabled when handling merge block fetch failure.") - mapOutputTracker. - unregisterMergeResult(shuffleId, reduceId, bmAddress, None) + if (mapStage.rdd.isBarrier()) { + // Mark all the map as broken in the map stage, to ensure retry all the tasks on + // resubmitted stage attempt. + // TODO: SPARK-35547: Clean all push-based shuffle metadata like merge enabled and + // TODO: finalized as we are clearing all the merge results. + mapOutputTracker.unregisterAllMapAndMergeOutput(shuffleId) + } else if (mapIndex != -1) { + // Mark the map whose fetch failed as broken in the map stage + mapOutputTracker.unregisterMapOutput(shuffleId, mapIndex, bmAddress) + if (pushBasedShuffleEnabled) { + // Possibly unregister the merge result , if the FetchFailed + // mapIndex is part of the merge result of + mapOutputTracker. + unregisterMergeResult(shuffleId, reduceId, bmAddress, Option(mapIndex)) + } + } else { + // Unregister the merge result of if there is a FetchFailed + // event and is not a MetaDataFetchException (signified by bmAddress being null) + if (bmAddress != null && + bmAddress.executorId.equals(BlockManagerId.SHUFFLE_MERGER_IDENTIFIER)) { + assert(pushBasedShuffleEnabled, "Push based shuffle expected to " + + "be enabled when handling merge block fetch failure.") + mapOutputTracker. + unregisterMergeResult(shuffleId, reduceId, bmAddress, None) + } } - } - if (failedStage.rdd.isBarrier()) { - failedStage match { - case failedMapStage: ShuffleMapStage => - // Mark all the map as broken in the map stage, to ensure retry all the tasks on - // resubmitted stage attempt. - mapOutputTracker.unregisterAllMapAndMergeOutput(failedMapStage.shuffleDep.shuffleId) - - case failedResultStage: ResultStage => - // Abort the failed result stage since we may have committed output for some - // partitions. - val reason = "Could not recover from a failed barrier ResultStage. Most recent " + - s"failure reason: $failureMessage" - abortStage(failedResultStage, reason, None) + if (failedStage.rdd.isBarrier()) { + failedStage match { + case failedMapStage: ShuffleMapStage => + // Mark all the map as broken in the map stage, to ensure retry all the tasks on + // resubmitted stage attempt. + mapOutputTracker.unregisterAllMapAndMergeOutput( + failedMapStage.shuffleDep.shuffleId) + + case failedResultStage: ResultStage => + // Abort the failed result stage since we may have committed output for some + // partitions. + val reason = "Could not recover from a failed barrier ResultStage. Most recent " + + s"failure reason: $failureMessage" + abortStage(failedResultStage, reason, None) + } } - } - if (shouldAbortStage) { - abortStage(failedStage, abortReason.get, None) - } else { // update failedStages and make sure a ResubmitFailedStages event is enqueued - // TODO: Cancel running tasks in the failed stage -- cf. SPARK-17064 - val noResubmitEnqueued = !failedStages.contains(failedStage) - failedStages += failedStage - failedStages += mapStage - if (noResubmitEnqueued) { - // For statically indeterminate stages, trigger rollback early (here and in - // submitMissingTasks) rather than deferring to task completion. This is more - // efficient because it clears shuffle outputs before the retry is submitted, - // ensuring findMissingPartitions() returns all partitions. - // - // For runtime detection (checksum mismatch), rollback is triggered at task - // completion when the mismatch is discovered. - // - // The `rollbackCurrentStage = true` parameter ensures the failed map stage is - // included in the cleanup: clearing its shuffle outputs, marking old task results - // to be ignored, and creating a new shuffle merge state for the upcoming retry. - if (mapStage.isStaticallyIndeterminate && - !mapStage.shuffleDep.checksumMismatchFullRetryEnabled) { - rollbackSucceedingStages(mapStage, rollbackCurrentStage = true) - } + if (shouldAbortStage) { + abortStage(failedStage, abortReason.get, None) + } else { // update failedStages and make sure a ResubmitFailedStages event is enqueued + // TODO: Cancel running tasks in the failed stage -- cf. SPARK-17064 + val noResubmitEnqueued = !failedStages.contains(failedStage) + failedStages += failedStage + failedStages += mapStage + if (noResubmitEnqueued) { + // For statically indeterminate stages, trigger rollback early (here and in + // submitMissingTasks) rather than deferring to task completion. This is more + // efficient because it clears shuffle outputs before the retry is submitted, + // ensuring findMissingPartitions() returns all partitions. + // + // For runtime detection (checksum mismatch), rollback is triggered at task + // completion when the mismatch is discovered. + // + // The `rollbackCurrentStage = true` parameter ensures the failed map stage is + // included in the cleanup: clearing its shuffle outputs, marking old task results + // to be ignored, and creating a new shuffle merge state for the upcoming retry. + if (mapStage.isStaticallyIndeterminate && + !mapStage.shuffleDep.checksumMismatchFullRetryEnabled) { + rollbackSucceedingStages(mapStage, rollbackCurrentStage = true) + } - // We expect one executor failure to trigger many FetchFailures in rapid succession, - // but all of those task failures can typically be handled by a single resubmission of - // the failed stage. We avoid flooding the scheduler's event queue with resubmit - // messages by checking whether a resubmit is already in the event queue for the - // failed stage. If there is already a resubmit enqueued for a different failed - // stage, that event would also be sufficient to handle the current failed stage, but - // producing a resubmit for each failed stage makes debugging and logging a little - // simpler while not producing an overwhelming number of scheduler events. - logInfo( - log"Resubmitting ${MDC(STAGE, mapStage)} " + - log"(${MDC(STAGE_NAME, mapStage.name)}) and ${MDC(FAILED_STAGE, failedStage)} " + - log"(${MDC(FAILED_STAGE_NAME, failedStage.name)}) due to fetch failure") - scheduleResubmit() + // We expect one executor failure to trigger many FetchFailures in rapid succession, + // but all of those task failures can typically be handled by a single resubmission + // of the failed stage. We avoid flooding the scheduler's event queue with resubmit + // messages by checking whether a resubmit is already in the event queue for the + // failed stage. If there is already a resubmit enqueued for a different failed + // stage, that event would also be sufficient to handle the current failed stage, + // but producing a resubmit for each failed stage makes debugging and logging a + // little simpler while not producing an overwhelming number of scheduler events. + logInfo( + log"Resubmitting ${MDC(STAGE, mapStage)} " + + log"(${MDC(STAGE_NAME, mapStage.name)}) and ${MDC(FAILED_STAGE, failedStage)} " + + log"(${MDC(FAILED_STAGE_NAME, failedStage.name)}) due to fetch failure") + scheduleResubmit() + } } } diff --git a/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala b/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala index da74803a10ebc..f36649828ec86 100644 --- a/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala +++ b/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala @@ -240,6 +240,46 @@ class PipelinedShuffleRoutingSuite extends SparkFunSuite with LocalSparkContext assert(SparkEnv.get.streamingShuffleOutputTracker.isEmpty) } + test("createShuffleMapStage fails loud for a pipelined dependency when no streaming tracker") { + // A pipelined shuffle is served solely by the StreamingShuffleOutputTracker. With a non- + // streaming incremental manager configured, that tracker is absent (verified above), so a job + // whose stage graph contains a PipelinedShuffleDependency cannot be scheduled. + // createShuffleMapStage must fail loud rather than silently register the shuffle in no tracker + // (which would strand a consumer with no writer locations). This guards the decoupling's + // invariant that a pipelined dependency implies a streaming tracker. + // local[4]: the pipelined group's up-front slot admission (producer 2 + result 2 = 4) must pass + // so that stage creation is actually reached -- otherwise the job is rejected for slots first. + sc = new SparkContext("local[4]", "test", newConf()) + assert(SparkEnv.get.streamingShuffleOutputTracker.isEmpty) + val producer = sc.parallelize(1 to 4, 2).map(x => (x, x)) + val pipelined = new PipelinedShuffleDependency[Int, Int, Int](producer, new HashPartitioner(2)) + // A minimal reduce-side RDD whose single dependency is the pipelined shuffle, so submitting an + // action forces createShuffleMapStage for the pipelined producer. + val consumer = new RDD[(Int, Int)](sc, Seq(pipelined)) { + override def compute(split: Partition, ctx: TaskContext): Iterator[(Int, Int)] = + Iterator.empty + override protected def getPartitions: Array[Partition] = + Array.tabulate(2)(i => new Partition { override def index: Int = i }) + } + val ex = intercept[Exception] { + consumer.count() + } + assert(findCause[IllegalStateException](ex).exists( + _.getMessage.contains("requires a StreamingShuffleOutputTracker")), + s"expected a fail-loud IllegalStateException about the missing tracker, got: $ex") + + // The fail-loud throw happens BEFORE createShuffleMapStage mutates stageIdToStage / + // shuffleIdToMapStage, so it leaves no partial scheduler state behind. If it left a half-created + // stage cached in shuffleIdToMapStage, a re-submission would reuse that stale stage instead of + // re-throwing. Re-submitting the same job must therefore throw the SAME fail-loud error again. + val ex2 = intercept[Exception] { + consumer.count() + } + assert(findCause[IllegalStateException](ex2).exists( + _.getMessage.contains("requires a StreamingShuffleOutputTracker")), + s"a re-submission must re-throw the fail-loud error (no leaked partial stage), got: $ex2") + } + test("spark.shuffle.manager.incremental resolves the same short aliases as the default manager") { // "sort" is a short alias the incremental slot must resolve (to SortShuffleManager) rather than // treat as a class name. SortShuffleManager is blocking, so it is rejected from the pipelined From ee7295983bacd750ff741607e76649799c79fbf2 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Wed, 29 Jul 2026 19:09:18 +0000 Subject: [PATCH 05/24] [SPARK-XXXXX][CORE] Fix scalastyle line length in the fail-loud routing test Co-authored-by: Isaac --- .../apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala b/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala index f36649828ec86..8e0f9e04c1bd8 100644 --- a/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala +++ b/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala @@ -269,9 +269,9 @@ class PipelinedShuffleRoutingSuite extends SparkFunSuite with LocalSparkContext s"expected a fail-loud IllegalStateException about the missing tracker, got: $ex") // The fail-loud throw happens BEFORE createShuffleMapStage mutates stageIdToStage / - // shuffleIdToMapStage, so it leaves no partial scheduler state behind. If it left a half-created - // stage cached in shuffleIdToMapStage, a re-submission would reuse that stale stage instead of - // re-throwing. Re-submitting the same job must therefore throw the SAME fail-loud error again. + // shuffleIdToMapStage, so it leaves no partial scheduler state behind. If it left a + // half-created stage cached in shuffleIdToMapStage, a re-submission would reuse that stale + // stage instead of re-throwing. Re-submitting the same job must therefore re-throw the error. val ex2 = intercept[Exception] { consumer.count() } From 3b24bb83a1980d0be8929384f08eaabfcadf6765 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Thu, 30 Jul 2026 06:00:30 +0000 Subject: [PATCH 06/24] [SPARK-XXXXX][CORE] Address review: fix comment line lengths and a stale availability comment - Reflow four 101-char comment lines in the FetchFailed group-abort branch and the unregisterOutputsOnFetchFailedExecutor scaladoc to <=100. - Correct ShuffleMapStage.numAvailableOutputs: a pipelined shuffle is not registered with the MapOutputTracker at all (it was leftover "registered but empty" wording that contradicted this file's own scaladoc). Comment-only; no behavior change. Co-authored-by: Isaac --- .../apache/spark/scheduler/DAGScheduler.scala | 16 ++++++++-------- .../apache/spark/scheduler/ShuffleMapStage.scala | 4 ++-- 2 files changed, 10 insertions(+), 10 deletions(-) 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 b64a9f5282aca..4e9ab23a21b22 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -3415,8 +3415,8 @@ private[spark] class DAGScheduler( (isPipelinedGroupMember(failedStage) || isPipelinedGroupMember(mapStage))) { // Failure is group-atomic for a pipelined group. The base scheduler handles a // FetchFailed by resubmitting just the map stage in isolation and recomputing serially, - // but a transient pipelined shuffle cannot be re-read and its members are co-scheduled, so - // a lone-stage resubmit is never valid and would deadlock the group. Abort the + // but a transient pipelined shuffle cannot be re-read and its members are co-scheduled, + // so a lone-stage resubmit is never valid and would deadlock the group. Abort the // whole group instead: aborting the failed stage tears down its running co-scheduled // members and fails the job, and the caller (e.g. the streaming batch loop) reruns the // batch from scratch. This is distinct from the maxTaskFailures=1 lever (which handles @@ -4111,12 +4111,12 @@ private[spark] class DAGScheduler( /** * On a FetchFailed, unregister the shuffle outputs of the executor (or its whole host) whose * fetch failed, treating the FetchFailed as authoritative evidence that its shuffle data is gone. - * Extracted from the base FetchFailed handler so the pipelined-group-abort branch can also run it: - * aborting the group fails only this job's stages, but a dead executor's REGULAR outputs must - * still be cleaned up for other/concurrent jobs (with an external shuffle service, an ExecutorLost - * does not clean them, so FetchFailed is the only proactive channel). No-op when `bmAddress` is - * null. Safe for a pipelined shuffle: it registers no map outputs in the tracker, so this can only - * strip regular/durable outputs. + * Extracted from the base FetchFailed handler so the pipelined-group-abort branch can also run + * it: aborting the group fails only this job's stages, but a dead executor's REGULAR outputs must + * still be cleaned up for other/concurrent jobs (with an external shuffle service, an + * ExecutorLost does not clean them, so FetchFailed is the only proactive channel). No-op when + * `bmAddress` is null. Safe for a pipelined shuffle: it registers no map outputs in the tracker, + * so this can only strip regular/durable outputs. */ private def unregisterOutputsOnFetchFailedExecutor( bmAddress: BlockManagerId, task: Task[_]): Unit = { diff --git a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala index 135e5e5a4bf2c..11cec94f2b311 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala @@ -109,8 +109,8 @@ private[spark] class ShuffleMapStage( * When this reaches [[numPartitions]], this map stage is ready. */ def numAvailableOutputs: Int = { - // A pipelined shuffle's outputs are never populated in the MapOutputTracker (its entry stays - // empty; see pipelinedCompletedPartitions); read its locally-tracked, monotonic completed set. + // A pipelined shuffle is not registered with the MapOutputTracker at all (see + // pipelinedCompletedPartitions); read its locally-tracked, monotonic completed set instead. if (isPipelined) pipelinedCompletedPartitions.size else mapOutputTrackerMaster.getNumAvailableOutputs(shuffleDep.shuffleId) } From 1f2ae87b0de140e98a86ab7e78f59b4a5fa8c5b0 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Thu, 30 Jul 2026 16:09:30 +0000 Subject: [PATCH 07/24] [SPARK-XXXXX][CORE] Address post-merge review comments from #57361 cloud-fan's nits on the merged predecessor PR, applied here per his request to address them in the follow-up: - DAGScheduler: reword the push-based-shuffle-merge rejection comment ("as the" -> "on a" pipelined shuffle) and the slot-check-disabled warning comment ("Keeps" -> "Limits the warning to once"). - TaskSetManagerSuite: the pipelined-decommission test comment said the shuffle is "NEVER registered in the MapOutputTracker". With the MapOutputTracker now fully decoupled that is literally true, but state it precisely: registered only with the StreamingShuffleOutputTracker, never the MapOutputTracker. Comment-only; no behavior change. Co-authored-by: Isaac --- .../main/scala/org/apache/spark/scheduler/DAGScheduler.scala | 4 ++-- .../org/apache/spark/scheduler/TaskSetManagerSuite.scala | 5 +++-- 2 files changed, 5 insertions(+), 4 deletions(-) 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 4e9ab23a21b22..d333b32c94e25 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -188,7 +188,7 @@ private[spark] class DAGScheduler( private[scheduler] val dependentStageMap = new HashMap[Stage, DependentStageInfo] // Whether we have already logged that the pipelined-group slot check is disabled. Only the - // (single-threaded) event loop touches this, so a plain var is safe. Keeps the warning to once + // (single-threaded) event loop touches this, so a plain var is safe. Limits the warning to once // per scheduler rather than once per submitted batch job. private var warnedPipelinedSlotCheckDisabled = false @@ -809,7 +809,7 @@ private[spark] class DAGScheduler( if (shuffleDep.checksumMismatchFullRetryEnabled) { throw pipelinedUnsupportedError("checksum-mismatch full retry with a pipelined shuffle") } - // Push-based shuffle merge as the pipelined shuffle: exposes output only after a + // Push-based shuffle merge on a pipelined shuffle: exposes output only after a // post-completion finalize step, the opposite of incremental reads. A // PipelinedShuffleDependency disables merge in its constructor, so this is a defensive // backstop against that being bypassed. diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala index 4c7e8efc4eaba..5fc7c842c54b0 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala @@ -3139,8 +3139,9 @@ class TaskSetManagerSuite // The TaskSetManager.executorLost "Resubmitted" loop re-enqueues an already-successful // ShuffleMapTask when its executor is lost and the map output looks gone -- for a decommission, // it looks gone whenever MapOutputTracker.getMapOutputLocation returns None. A pipelined - // shuffle is NEVER registered in the MapOutputTracker, so getMapOutputLocation is always None - // and this loop would resubmit the lone producer task -- the exact streaming-writer hang. This + // shuffle is registered only with the StreamingShuffleOutputTracker, never the + // MapOutputTracker, so getMapOutputLocation is always None and this loop would resubmit the + // lone producer task -- the exact streaming-writer hang. This // loop bypasses handleFailedTask, so the group-atomic abort would never see it. The guard // (!taskSet.isPipelined on maybeShuffleMapOutputLoss) must keep a pipelined set out of the // loop. From 215133c8b8b2a5c63ecd51222377dc8ad384e6ab Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Thu, 30 Jul 2026 23:01:09 +0000 Subject: [PATCH 08/24] [SPARK-XXXXX][CORE] Select the output tracker by shuffle-dependency type via a common interface Introduce a small ShuffleOutputTrackerMaster trait (registerShuffle / containsShuffle / unregisterShuffle) implemented by both MapOutputTrackerMaster and StreamingShuffleOutputTrackerMaster. DAGScheduler.createShuffleMapStage now resolves the owning tracker by dependency type through a private outputTrackerMaster(shuffleDep) helper (which isolates the fail-loud-if-no-streaming-tracker check) and registers through the common interface, replacing the awkward Option[StreamingShuffleOutputTrackerMaster] + isInstanceOf/asInstanceOf and the type-specific registration branches. Behavior is unchanged. Co-authored-by: Isaac --- .../org/apache/spark/MapOutputTracker.scala | 10 +- .../spark/StreamingShuffleOutputTracker.scala | 25 ++++- .../apache/spark/scheduler/DAGScheduler.scala | 103 ++++++++---------- 3 files changed, 73 insertions(+), 65 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/MapOutputTracker.scala b/core/src/main/scala/org/apache/spark/MapOutputTracker.scala index 742baf40ecbf0..933da94772725 100644 --- a/core/src/main/scala/org/apache/spark/MapOutputTracker.scala +++ b/core/src/main/scala/org/apache/spark/MapOutputTracker.scala @@ -799,7 +799,7 @@ private[spark] class MapOutputTrackerMaster( conf: SparkConf, private[spark] val broadcastManager: BroadcastManager, private[spark] val isLocal: Boolean) - extends MapOutputTracker(conf) { + extends MapOutputTracker(conf) with ShuffleOutputTrackerMaster { // The size at which we use Broadcast to send the map output statuses to the executors private val minSizeForBroadcast = conf.get(SHUFFLE_MAPOUTPUT_MIN_SIZE_FOR_BROADCAST).toInt @@ -943,6 +943,10 @@ private[spark] class MapOutputTrackerMaster( } } + // ShuffleOutputTrackerMaster: a regular shuffle has no per-job registration, so jobId is ignored. + override def registerShuffle(shuffleId: Int, numMaps: Int, numReduces: Int, jobId: Int): Unit = + registerShuffle(shuffleId, numMaps, numReduces) + def updateMapOutput(shuffleId: Int, mapId: Long, bmAddress: BlockManagerId): Unit = { shuffleStatuses.get(shuffleId) match { case Some(shuffleStatus) => @@ -1030,7 +1034,7 @@ private[spark] class MapOutputTrackerMaster( } /** Unregister shuffle data */ - def unregisterShuffle(shuffleId: Int): Unit = { + override def unregisterShuffle(shuffleId: Int): Unit = { shuffleStatuses.remove(shuffleId).foreach { shuffleStatus => shuffleStatus.invalidateSerializedMapOutputStatusCache() shuffleStatus.invalidateSerializedMergeOutputStatusCache() @@ -1077,7 +1081,7 @@ private[spark] class MapOutputTrackerMaster( } /** Check if the given shuffle is being tracked */ - def containsShuffle(shuffleId: Int): Boolean = shuffleStatuses.contains(shuffleId) + override def containsShuffle(shuffleId: Int): Boolean = shuffleStatuses.contains(shuffleId) def getNumAvailableOutputs(shuffleId: Int): Int = { shuffleStatuses.get(shuffleId).map(_.numAvailableMapOutputs).getOrElse(0) diff --git a/core/src/main/scala/org/apache/spark/StreamingShuffleOutputTracker.scala b/core/src/main/scala/org/apache/spark/StreamingShuffleOutputTracker.scala index 9d6fc4f1936c8..df4e76b017dbc 100644 --- a/core/src/main/scala/org/apache/spark/StreamingShuffleOutputTracker.scala +++ b/core/src/main/scala/org/apache/spark/StreamingShuffleOutputTracker.scala @@ -28,6 +28,23 @@ import org.apache.spark.internal.config.SHUFFLE_MAPOUTPUT_DISPATCHER_NUM_THREADS import org.apache.spark.rpc.{RpcCallContext, RpcEndpoint, RpcEndpointRef, RpcEnv} import org.apache.spark.util.ThreadUtils +/** + * The driver-side registry for a shuffle's output. A regular shuffle is served by the + * `MapOutputTrackerMaster`, a pipelined (streaming) shuffle by the + * `StreamingShuffleOutputTrackerMaster` -- split by dependency type, with no overlap. This is the + * small common surface the DAGScheduler and ContextCleaner drive polymorphically, so they select + * the right tracker by shuffle-dependency type (see `DAGScheduler.outputTrackerMaster`) rather than + * special-casing pipelined shuffles at each call site. + */ +private[spark] trait ShuffleOutputTrackerMaster { + /** Register a shuffle so its outputs can be tracked. `jobId` is used by the streaming tracker. */ + def registerShuffle(shuffleId: Int, numMaps: Int, numReduces: Int, jobId: Int): Unit + /** Whether the given shuffle is registered with this tracker. */ + def containsShuffle(shuffleId: Int): Boolean + /** Unregister a shuffle and release its tracked state. */ + def unregisterShuffle(shuffleId: Int): Unit +} + private[spark] sealed trait StreamingShuffleTaskLocationTrackerMessage private[spark] case class UpdateStreamingShuffleTaskLocation( @@ -196,7 +213,7 @@ private[spark] abstract class StreamingShuffleOutputTracker(conf: SparkConf) ext private[spark] case class StreamingShuffleInfo(numMaps: Int, numReduces: Int, jobId: Int) private[spark] class StreamingShuffleOutputTrackerMaster(conf: SparkConf) - extends StreamingShuffleOutputTracker(conf) { + extends StreamingShuffleOutputTracker(conf) with ShuffleOutputTrackerMaster { // map that stores task location information organized in the following fashion // shuffle id -> {mapId -> location} @@ -220,7 +237,7 @@ private[spark] class StreamingShuffleOutputTrackerMaster(conf: SparkConf) pool } - def registerShuffle(shuffleId: Int, numMaps: Int, numReduces: Int, jobId: Int): Unit = { + override def registerShuffle(shuffleId: Int, numMaps: Int, numReduces: Int, jobId: Int): Unit = { logInfo(log"Registering shuffleId ${MDC(LogKeys.SHUFFLE_ID, shuffleId)} with ${ MDC(LogKeys.NUM_MAPPERS, numMaps)} mappers and ${ MDC(LogKeys.NUM_REDUCERS, numReduces)} reducers") @@ -230,9 +247,7 @@ private[spark] class StreamingShuffleOutputTrackerMaster(conf: SparkConf) } } - // Whether the given shuffle is registered. ContextCleaner uses this on the driver to skip - // regular (non-pipelined) shuffles, which are never registered here, before unregistering. - def containsShuffle(shuffleId: Int): Boolean = shuffleInfos.containsKey(shuffleId) + override def containsShuffle(shuffleId: Int): Boolean = shuffleInfos.containsKey(shuffleId) // for testing purposes private[spark] def getShuffleInfo(shuffleId: Int): Option[StreamingShuffleInfo] = { 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 d333b32c94e25..05c142558a435 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -692,24 +692,11 @@ private[spark] class DAGScheduler( checkBarrierStageWithNumSlots(rdd, resourceProfile) checkBarrierStageWithRDDChainPattern(rdd, rdd.getNumPartitions) checkPipelinedProducerSupported(shuffleDep) - // A pipelined shuffle is served solely by the StreamingShuffleOutputTracker, so that tracker - // MUST exist for one -- it is created whenever a streaming-capable shuffle manager is set up - // (see SparkEnv.initializeStreamingShuffleOutputTracker), a prerequisite for producing a - // PipelinedShuffleDependency. Resolve it up front (BEFORE any stage-map mutation below), so a - // fail-loud on the misconfiguration leaves no partial scheduler state -- and so a pipelined - // shuffle is never silently registered in no tracker (a consumer would then find no writer - // locations). The reader enforces the same invariant (see StreamingShuffleReader). None for a - // regular shuffle, which is registered with the MapOutputTracker instead. - val streamingTrackerForPipelined: Option[StreamingShuffleOutputTrackerMaster] = - if (shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) { - Some(sc.env.streamingShuffleOutputTracker - .getOrElse(throw new IllegalStateException( - s"A pipelined shuffle (id ${shuffleDep.shuffleId}) requires a " + - "StreamingShuffleOutputTracker, but none is configured")) - .asInstanceOf[StreamingShuffleOutputTrackerMaster]) - } else { - None - } + // Resolve the tracker that owns this shuffle's output, by dependency type (see + // outputTrackerMaster). Resolve it up front, BEFORE any stage-map mutation below, so that a + // fail-loud on a misconfigured pipelined shuffle (no StreamingShuffleOutputTracker) leaves no + // partial scheduler state and a re-submit re-throws the same error. + val outputTracker = outputTrackerMaster(shuffleDep) val numTasks = rdd.partitions.length val parents = getOrCreateParentStages(shuffleDeps, jobId) val id = nextStageId.getAndIncrement() @@ -721,49 +708,51 @@ private[spark] class DAGScheduler( shuffleIdToMapStage(shuffleDep.shuffleId) = stage updateJobIdStageIdMaps(jobId, stage) - // A pipelined shuffle and a regular shuffle live in DIFFERENT trackers -- the split is by - // dependency type, with no overlap. A pipelined shuffle produces no durable, addressable map - // output; its producer/consumer task-location routing lives entirely in the - // StreamingShuffleOutputTracker (resolved fail-loud above), and it is never registered with the - // MapOutputTracker (nothing reads such an entry -- availability is tracked on the stage via - // pipelinedCompletedPartitions, and the reader locates writers through the streaming tracker). - // The MapOutputTracker.getStatistics paths, which WOULD throw ShuffleStatusNotFoundException on - // the absent entry, are all unreachable for a pipelined dependency: markMapStageJobsAsFinished - // only calls it when mapStageJobs is non-empty, but handleMapStageSubmitted rejects a pipelined - // dependency before addActiveJob (its sole populator) ever runs, so a pipelined stage's - // mapStageJobs is always empty; and the getStatistics in checkAndScheduleShuffleMergeFinalize - // is on the push-based-shuffle merge path, which a pipelined dependency rejects up front - // (shuffleMergeEnabled in checkPipelinedProducerSupported). A regular shuffle registers with - // the MapOutputTracker as before. - streamingTrackerForPipelined match { - case Some(streamingTracker) => - // Register the producer/consumer fan-in before any consumer task asks for a writer loc. - // Self-guarded on the tracker's own state (createShuffleMapStage runs once per shuffleId - // via getOrCreateShuffleMapStage, but guard defensively against a re-entry). - if (!streamingTracker.containsShuffle(shuffleDep.shuffleId)) { - logInfo(log"Registering RDD ${MDC(RDD_ID, rdd.id)} " + - log"(${MDC(CREATION_SITE, rdd.getCreationSite)}) as input to pipelined " + - log"shuffle ${MDC(SHUFFLE_ID, shuffleDep.shuffleId)}") - streamingTracker.registerShuffle( - shuffleDep.shuffleId, - rdd.partitions.length, - shuffleDep.partitioner.numPartitions, - jobId) - } - case None if !mapOutputTracker.containsShuffle(shuffleDep.shuffleId) => - // Kind of ugly: need to register RDDs with the cache and map output tracker here - // since we can't do it in the RDD constructor because # of partitions is unknown - logInfo(log"Registering RDD ${MDC(RDD_ID, rdd.id)} " + - log"(${MDC(CREATION_SITE, rdd.getCreationSite)}) as input to " + - log"shuffle ${MDC(SHUFFLE_ID, shuffleDep.shuffleId)}") - mapOutputTracker.registerShuffle(shuffleDep.shuffleId, rdd.partitions.length, - shuffleDep.partitioner.numPartitions) - - case _ => // regular shuffle already registered with the MapOutputTracker; nothing to do + // Register the shuffle with its own tracker (a pipelined shuffle in the + // StreamingShuffleOutputTracker, a regular one in the MapOutputTracker -- split by dependency + // type, no overlap). Self-guarded on the tracker's own membership: createShuffleMapStage runs + // once per shuffleId via getOrCreateShuffleMapStage, but guard defensively against re-entry. + // (A pipelined shuffle is never registered with the MapOutputTracker; its availability is + // tracked on the stage via pipelinedCompletedPartitions and its writers are located through the + // StreamingShuffleOutputTracker. jobId is used only by the streaming tracker.) The + // MapOutputTracker.getStatistics paths, which WOULD throw ShuffleStatusNotFoundException on the + // absent entry, are unreachable for a pipelined dependency: markMapStageJobsAsFinished calls it + // only when mapStageJobs is non-empty, but handleMapStageSubmitted rejects a pipelined dep + // before addActiveJob (its sole populator) runs, so a pipelined stage's mapStageJobs is always + // empty; and checkAndScheduleShuffleMergeFinalize's getStatistics is on the push-based-merge + // path, which a pipelined dependency rejects up front (checkPipelinedProducerSupported). + if (!outputTracker.containsShuffle(shuffleDep.shuffleId)) { + logInfo(log"Registering RDD ${MDC(RDD_ID, rdd.id)} " + + log"(${MDC(CREATION_SITE, rdd.getCreationSite)}) as input to " + + log"shuffle ${MDC(SHUFFLE_ID, shuffleDep.shuffleId)}") + outputTracker.registerShuffle(shuffleDep.shuffleId, rdd.partitions.length, + shuffleDep.partitioner.numPartitions, jobId) } stage } + /** + * The driver-side output tracker that owns a shuffle's outputs, selected by dependency type: the + * StreamingShuffleOutputTracker for a pipelined shuffle, the MapOutputTracker otherwise. The two + * are split with no overlap. A pipelined shuffle REQUIRES a StreamingShuffleOutputTracker + * (created with a streaming-capable shuffle manager, see + * SparkEnv.initializeStreamingShuffleOutputTracker), so fail loud if one is not configured rather + * than silently register it nowhere (a consumer would then find no writer locations; the reader + * enforces the same invariant, see StreamingShuffleReader). + */ + private def outputTrackerMaster( + shuffleDep: ShuffleDependency[_, _, _]): ShuffleOutputTrackerMaster = { + if (shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) { + sc.env.streamingShuffleOutputTracker + .getOrElse(throw new IllegalStateException( + s"A pipelined shuffle (id ${shuffleDep.shuffleId}) requires a " + + "StreamingShuffleOutputTracker, but none is configured")) + .asInstanceOf[StreamingShuffleOutputTrackerMaster] + } else { + mapOutputTracker + } + } + private def pipelinedUnsupportedError(reason: String): PipelinedShuffleUnsupportedException = new PipelinedShuffleUnsupportedException(reason) From 142850bddd98dcb41f5a5ca9d129157f8dad1b6c Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Thu, 30 Jul 2026 23:01:24 +0000 Subject: [PATCH 09/24] [SPARK-XXXXX][CORE] Add leak-coverage tests for pipelined-shuffle cleanup Close two gaps in the pipelined-shuffle leak coverage: - ContextCleanerSuite: a GC-triggered end-to-end cleanup test. A pipelined shuffle's only cleanup channel is the ContextCleaner weak-reference path fired when its PipelinedShuffleDependency is garbage-collected; the existing tests call doCleanupShuffle directly. The new test builds a real dependency, asserts cleanup does NOT fire while it is strongly referenced, then dereferences it and asserts GC drives the streaming-tracker unregister. - DAGSchedulerSuite: a fan-in (two-producer) deferral leak test that drives a producer failure to full job teardown and asserts all scheduler state (including the two-parent deferral) is emptied. Co-authored-by: Isaac --- .../apache/spark/ContextCleanerSuite.scala | 43 +++++++++++++++++++ .../spark/scheduler/DAGSchedulerSuite.scala | 40 +++++++++++++++++ 2 files changed, 83 insertions(+) diff --git a/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala b/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala index e8f33b7803da9..0bd0ee103a190 100644 --- a/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala +++ b/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala @@ -221,6 +221,49 @@ class ContextCleanerSuite extends ContextCleanerSuiteBase { postGCTester.assertCleanup() } + test("automatically cleanup pipelined shuffle from the streaming output tracker") { + // The load-bearing leak guarantee for a pipelined (streaming) shuffle: it is registered ONLY + // with the driver-only StreamingShuffleOutputTracker (never the MapOutputTracker), and its only + // cleanup channel is the ContextCleaner weak-reference path fired when the + // PipelinedShuffleDependency is garbage-collected (registerShuffleForCleanup in the + // ShuffleDependency constructor). If that GC path did not reach the streaming tracker, + // the tracker would grow without bound across Real-Time Mode micro-batches. The sibling + // doCleanupShuffle tests call cleanup directly; this one proves the end-to-end + // GC -> ContextCleaner -> streaming-tracker unregister link. + val streamingTracker = sc.env.streamingShuffleOutputTracker.get + .asInstanceOf[StreamingShuffleOutputTrackerMaster] + val mapOutputTracker = sc.env.mapOutputTracker.asInstanceOf[MapOutputTrackerMaster] + + // Build a real PipelinedShuffleDependency (registers itself for cleanup on construction) and + // register its shuffleId in the streaming tracker exactly as DAGScheduler.createShuffleMapStage + // does for a pipelined shuffle (nothing is put in the MapOutputTracker). + var pipelinedDep: PipelinedShuffleDependency[Int, Int, Int] = + new PipelinedShuffleDependency(newPairRDD(), new HashPartitioner(2)) + val shuffleId = pipelinedDep.shuffleId + streamingTracker.registerShuffle(shuffleId, numMaps = 2, numReduces = 2, jobId = 0) + assert(streamingTracker.containsShuffle(shuffleId)) + assert(!mapOutputTracker.containsShuffle(shuffleId), + "a pipelined shuffle must not be registered with the MapOutputTracker") + + // A strong reference to the dependency must prevent GC-triggered cleanup. + runGC() + intercept[Exception] { + eventually(timeout(1.second), interval(100.milliseconds)) { + assert(!streamingTracker.containsShuffle(shuffleId), + "cleanup must NOT fire while the dependency is strongly referenced") + } + } + + // Dereference the dependency; GC must then drive ContextCleaner to unregister it from the + // streaming tracker (and only there -- it was never in the MapOutputTracker). + pipelinedDep = null + runGC() + eventually(timeout(10.seconds), interval(100.milliseconds)) { + assert(!streamingTracker.containsShuffle(shuffleId), + "GC of the PipelinedShuffleDependency must unregister it from the streaming tracker") + } + } + test("automatically cleanup broadcast") { var broadcast = newBroadcast() 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 05ba990e2d638..a0428d8d8bd49 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -7078,6 +7078,46 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti s"a dropped fan-in consumer must not replay when a surviving producer succeeds; got $results") } + test("pipelined shuffle: a fan-in group leaks no scheduler state when a producer fails and the " + + "job is torn down") { + // Leak check for the multi-producer (fan-in) deferral: a consumer deferred against TWO + // producers buffers its completions; when one producer fails and the job is torn down, ALL + // scheduler state -- the consumer's dependentStageMap deferral keyed on both producers -- + // must be cleaned up. The sibling "dropped as soon as one producer fails" test stops at the + // drop-vs-replay outcome; this one drives the failure to full job teardown and asserts no leak. + val producerA = new MyRDD(sc, 2, Nil) + val psdA = new PipelinedShuffleDependency(producerA, new HashPartitioner(2)) + val producerB = new MyRDD(sc, 2, Nil) + val psdB = new PipelinedShuffleDependency(producerB, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(psdA, psdB), tracker = mapOutputTracker) + val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = results.put(index, result) + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(consumerRdd, Array(0, 1), listener = failListener) + assert(taskSets.size === 3, s"A, B and the consumer must be co-scheduled; got ${taskSets.size}") + val consumerTaskSet = + taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd).get + val consumerStage = scheduler.stageIdToStage(consumerTaskSet.stageId) + val producerATaskSet = + taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq producerA).get + + // Consumer finishes early -> buffered (deferred) against BOTH producers. + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + assert(results.isEmpty, "consumer completions should be buffered while its producers run") + assert(scheduler.dependentStageMap.get(consumerStage).exists(_.parents.size == 2), + "the consumer must be deferred against both producers") + + // Producer A fails its whole task set -> the job fails and the group is torn down. + failed(producerATaskSet, "producer A blew up") + assert(failure.get() != null, "the job must fail when a fan-in producer fails") + assert(results.isEmpty, "buffered consumer successes must be dropped, not applied, on failure") + sc.listenerBus.waitUntilEmpty(10000) + // The whole point: no leak. Every scheduler map (including the two-parent deferral) empties. + assertDataStructuresEmpty() + } + test("pipelined shuffle: a multi-producer consumer's buffered successes are dropped when one " + "producer fails (no stale replay via a sibling)") { // Group-atomic drop across MULTIPLE producers. Consumer C reads TWO pipelined producers P1 and From e6fedd121087471b3948d2ecb6cc919831212307 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Thu, 30 Jul 2026 23:42:21 +0000 Subject: [PATCH 10/24] [SPARK-XXXXX][CORE] Balance shuffleDriverComponents.removeShuffle for a pipelined shuffle on cleanup The ShuffleDependency constructor calls shuffleDriverComponents.registerShuffle for EVERY shuffle, including a pipelined one, but the pipelined cleanup branch in ContextCleaner.doCleanupShuffle only unregistered from the StreamingShuffleOutputTracker -- it skipped the matching removeShuffle. Benign for the default (no-op registerShuffle) impl, but a custom ShuffleDriverComponents tracking per-shuffle driver state would leak it across micro-batches. Call removeShuffle in the pipelined branch too, matching the regular branch's ordering. (Found by Isaac review.) Co-authored-by: Isaac --- .../main/scala/org/apache/spark/ContextCleaner.scala | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/ContextCleaner.scala b/core/src/main/scala/org/apache/spark/ContextCleaner.scala index a01a480147d72..6c1b49157cc01 100644 --- a/core/src/main/scala/org/apache/spark/ContextCleaner.scala +++ b/core/src/main/scala/org/apache/spark/ContextCleaner.scala @@ -249,10 +249,15 @@ private[spark] class ContextCleaner( listeners.asScala.foreach(_.shuffleCleaned(shuffleId)) logDebug("Cleaned shuffle " + shuffleId) } else if (streamingShuffleOutputTrackerMaster.exists(_.containsShuffle(shuffleId))) { - // A pipelined shuffle's state is driver-only (the worker tracker caches nothing), so - // cleanup is a direct driver-side unregister. No RemoveShuffle RPC is needed: a streaming - // shuffle has no durable, block-manager-served files to remove. + // A pipelined shuffle's output state is driver-only (the worker tracker caches nothing), so + // it is unregistered directly from the StreamingShuffleOutputTracker rather than the + // MapOutputTracker. Still call shuffleDriverComponents.removeShuffle to balance the + // registerShuffle that the ShuffleDependency constructor issues for EVERY shuffle (incl. a + // pipelined one): a custom components impl may track per-shuffle driver state that would + // otherwise leak. For the default impl this is just a RemoveShuffle RPC that finds no + // blocks (a pipelined shuffle has none), matching the regular branch's ordering. logDebug("Cleaning pipelined shuffle " + shuffleId) + shuffleDriverComponents.removeShuffle(shuffleId, blocking) streamingShuffleOutputTrackerMaster.foreach(_.unregisterShuffle(shuffleId)) listeners.asScala.foreach(_.shuffleCleaned(shuffleId)) logDebug("Cleaned pipelined shuffle " + shuffleId) From 4d6a5b25ec37854640501023e6486a2638774804 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Wed, 22 Jul 2026 04:43:06 +0000 Subject: [PATCH 11/24] [SPARK-XXXXX][SQL] Let ShuffleExchangeExec build a pipelined shuffle dependency Add a `pipelined` flag to ShuffleExchangeExec. When set, prepareShuffleDependency constructs a PipelinedShuffleDependency instead of a regular ShuffleDependency, so the DAGScheduler co-schedules the exchange's producer and consumer stages and routes the shuffle to the streaming shuffle manager. The flag defaults to false, so every existing shuffle is unchanged. Modeling this as a node field (rather than out-of-band metadata) keeps the pipelined/non-pipelined distinction part of the exchange's identity, so it participates in exchange-reuse and canonicalization -- correct for a transient shuffle that must not be reused as a durable one. Adding the field shifts ShuffleExchangeExec's positional arity, so update the pattern matches that destructure it positionally (EnsureRequirements, AQEUtils, InsertSortForLimitAndOffset and several test suites) to account for the new field; EnsureRequirements' repartition rewrite now uses copy() so the flag is preserved. Co-authored-by: Isaac --- .../InsertSortForLimitAndOffset.scala | 2 +- .../sql/execution/adaptive/AQEUtils.scala | 2 +- .../exchange/EnsureRequirements.scala | 6 +- .../exchange/ShuffleExchangeExec.scala | 42 +++++-- .../sql/DataFrameWindowFunctionsSuite.scala | 2 +- .../org/apache/spark/sql/DatasetSuite.scala | 2 +- .../spark/sql/execution/PlannerSuite.scala | 5 +- .../exchange/EnsureRequirementsSuite.scala | 110 +++++++++--------- .../sql/streaming/StreamingJoinSuite.scala | 4 +- 9 files changed, 97 insertions(+), 78 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffset.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffset.scala index aa29128cda7e0..88f0de2555f70 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffset.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffset.scala @@ -44,7 +44,7 @@ object InsertSortForLimitAndOffset extends Rule[SparkPlan] { _, // Should not match AQE shuffle stage because we only target un-submitted stages which // we can still rewrite the query plan. - s @ ShuffleExchangeExec(SinglePartition, child, _, _), + s @ ShuffleExchangeExec(SinglePartition, child, _, _, _), _) if child.logicalLink.isDefined => extractOrderingAndPropagateOrderingColumns(child) match { case Some((ordering, newChild)) => diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEUtils.scala index 578e0acd80525..cc3ca90739090 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEUtils.scala @@ -30,7 +30,7 @@ object AQEUtils { // Project/Filter/LocalSort/CollectMetrics. // Note: we only care about `HashPartitioning` as `EnsureRequirements` can only optimize out // user-specified repartition with `HashPartitioning`. - case ShuffleExchangeExec(h: HashPartitioning, _, shuffleOrigin, _) + case ShuffleExchangeExec(h: HashPartitioning, _, shuffleOrigin, _, _) if shuffleOrigin == REPARTITION_BY_COL || shuffleOrigin == REPARTITION_BY_NUM => val numPartitions = if (shuffleOrigin == REPARTITION_BY_NUM) { Some(h.numPartitions) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala index c632b3d841e61..8ca344e511946 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala @@ -272,8 +272,8 @@ case class EnsureRequirements( } child match { - case ShuffleExchangeExec(_, c, so, ps) => - ShuffleExchangeExec(newPartitioning, c, so, ps) + case s: ShuffleExchangeExec => + s.copy(outputPartitioning = newPartitioning) case gpe: GroupPartitionsExec => ShuffleExchangeExec(newPartitioning, gpe.child) case _ => ShuffleExchangeExec(newPartitioning, child) } @@ -896,7 +896,7 @@ case class EnsureRequirements( def apply(plan: SparkPlan): SparkPlan = { val newPlan = plan.transformUp { - case operator @ ShuffleExchangeExec(upper: HashPartitioning, child, shuffleOrigin, _) + case operator @ ShuffleExchangeExec(upper: HashPartitioning, child, shuffleOrigin, _, _) if optimizeOutRepartition && (shuffleOrigin == REPARTITION_BY_COL || shuffleOrigin == REPARTITION_BY_NUM) => def hasSemanticEqualPartitioning(partitioning: Partitioning): Boolean = { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala index dd829e697df61..92893a332ba08 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala @@ -191,7 +191,8 @@ case class ShuffleExchangeExec( override val outputPartitioning: Partitioning, child: SparkPlan, shuffleOrigin: ShuffleOrigin = ENSURE_REQUIREMENTS, - advisoryPartitionSize: Option[Long] = None) + advisoryPartitionSize: Option[Long] = None, + pipelined: Boolean = false) extends ShuffleExchangeLike { private lazy val writeMetrics = @@ -252,7 +253,8 @@ case class ShuffleExchangeExec( child.output, outputPartitioning, serializer, - writeMetrics) + writeMetrics, + pipelined) metrics("numPartitions").set(dep.partitioner.numPartitions) val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY) SQLMetrics.postDriverMetricUpdates( @@ -346,7 +348,8 @@ object ShuffleExchangeExec { outputAttributes: Seq[Attribute], newPartitioning: Partitioning, serializer: Serializer, - writeMetrics: Map[String, SQLMetric]) + writeMetrics: Map[String, SQLMetric], + pipelined: Boolean = false) : ShuffleDependency[Int, InternalRow, InternalRow] = { val part: Partitioner = newPartitioning match { case RoundRobinPartitioning(numPartitions) => new HashPartitioner(numPartitions) @@ -537,15 +540,30 @@ object ShuffleExchangeExec { } } val dependency = - new ShuffleDependency[Int, InternalRow, InternalRow]( - rddWithPartitionIds, - new PartitionIdPassthrough(part.numPartitions), - serializer, - shuffleWriterProcessor = createShuffleWriteProcessor(writeMetrics), - rowBasedChecksums = UnsafeRowChecksum.createUnsafeRowChecksums(checksumSize), - _checksumMismatchFullRetryEnabled = SQLConf.get.shuffleChecksumMismatchFullRetryEnabled, - checksumMismatchQueryLevelRollbackEnabled = - SQLConf.get.shuffleChecksumMismatchQueryLevelRollbackEnabled) + if (pipelined) { + // A pipelined shuffle is transient and incrementally readable: the DAGScheduler + // co-schedules its producer and consumer stages instead of materializing the shuffle + // first. The PipelinedShuffleDependency type is the entire opt-in -- routing to the + // streaming shuffle manager and pipelined-group co-scheduling both follow from it. The + // checksum-mismatch retry knobs are intentionally not carried over: a transient shuffle is + // never recomputed, so PipelinedShuffleDependency does not expose them (they stay off). + new PipelinedShuffleDependency[Int, InternalRow, InternalRow]( + rddWithPartitionIds, + new PartitionIdPassthrough(part.numPartitions), + serializer, + shuffleWriterProcessor = createShuffleWriteProcessor(writeMetrics), + rowBasedChecksums = UnsafeRowChecksum.createUnsafeRowChecksums(checksumSize)) + } else { + new ShuffleDependency[Int, InternalRow, InternalRow]( + rddWithPartitionIds, + new PartitionIdPassthrough(part.numPartitions), + serializer, + shuffleWriterProcessor = createShuffleWriteProcessor(writeMetrics), + rowBasedChecksums = UnsafeRowChecksum.createUnsafeRowChecksums(checksumSize), + _checksumMismatchFullRetryEnabled = SQLConf.get.shuffleChecksumMismatchFullRetryEnabled, + checksumMismatchQueryLevelRollbackEnabled = + SQLConf.get.shuffleChecksumMismatchQueryLevelRollbackEnabled) + } dependency } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala index c1d1f13915a6e..136dc57fdfcfa 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala @@ -1360,7 +1360,7 @@ class DataFrameWindowFunctionsSuite extends SharedSparkSession def isShuffleExecByRequirement( plan: ShuffleExchangeExec, desiredClusterColumns: Seq[String]): Boolean = plan match { - case ShuffleExchangeExec(op: HashPartitioning, _, ENSURE_REQUIREMENTS, _) => + case ShuffleExchangeExec(op: HashPartitioning, _, ENSURE_REQUIREMENTS, _, _) => partitionExpressionsColumns(op.expressions) === desiredClusterColumns case _ => false } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala index 3b28cae31a134..096eee137b320 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala @@ -1942,7 +1942,7 @@ class DatasetSuite extends SharedSparkSession val agg = cp.groupBy($"id" % 2).agg(count($"id")) agg.queryExecution.executedPlan.collectFirst { - case ShuffleExchangeExec(_, _: RDDScanExec, _, _) => + case ShuffleExchangeExec(_, _: RDDScanExec, _, _, _) => case BroadcastExchangeExec(_, _: RDDScanExec) => }.foreach { _ => fail( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala index 71192e92af0d0..41b74c6d4ff21 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala @@ -733,10 +733,11 @@ class PlannerSuite extends SharedSparkSession with AdaptiveSparkPlanHelper { outputPlan match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, - ShuffleExchangeExec(HashPartitioning(leftPartitioningExpressions, _), _, _, _), _), + ShuffleExchangeExec(HashPartitioning(leftPartitioningExpressions, _), + _, _, _, _), _), SortExec(_, _, ShuffleExchangeExec(HashPartitioning(rightPartitioningExpressions, _), - _, _, _), _), _) => + _, _, _, _), _), _) => assert(leftKeys === smjExec.leftKeys) assert(rightKeys === smjExec.rightKeys) assert(leftKeys === leftPartitioningExpressions) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala index 17d00ec055e07..95085b71bc251 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala @@ -60,7 +60,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec1) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _), _), - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB)) assert(rightKeys === Seq(exprB, exprA)) case other => fail(other.toString) @@ -71,7 +71,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprB :: exprA :: Nil, Inner, None, plan2, plan1) EnsureRequirements.apply(smjExec2) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _), _), _) => assert(leftKeys === Seq(exprB, exprA)) assert(rightKeys === Seq(exprA, exprB)) @@ -84,7 +84,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprD :: exprC :: Nil, exprB :: exprA :: Nil, Inner, None, plan1, plan1) EnsureRequirements.apply(smjExec3) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _), _), _) => assert(leftKeys === Seq(exprC, exprD)) assert(rightKeys === Seq(exprA, exprB)) @@ -126,8 +126,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { ) EnsureRequirements.apply(smjExec2) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprC, exprB, exprD)) assert(rightKeys === Seq(exprD, exprA, exprC)) case other => fail(other.toString) @@ -145,7 +145,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprB :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec1) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), _) => assert(leftKeys === Seq(exprB, exprA)) assert(rightKeys === Seq(exprB, exprC)) @@ -159,7 +159,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprB :: Nil, Inner, None, plan1, plan3) EnsureRequirements.apply(smjExec2) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _), _), _) => assert(leftKeys === Seq(exprB, exprA)) assert(rightKeys === Seq(exprB, exprC)) @@ -173,7 +173,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec3) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _), _), - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprB, exprC)) assert(rightKeys === Seq(exprB, exprA)) case other => fail(other.toString) @@ -319,7 +319,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB)) assert(rightKeys === Seq(exprC, exprD)) assert(p.expressions == Seq(exprC)) @@ -335,7 +335,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB)) assert(rightKeys === Seq(exprC, exprD)) @@ -353,7 +353,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB)) assert(rightKeys === Seq(exprC, exprD)) assert(p.expressions == Seq(exprC)) @@ -372,7 +372,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB, exprB)) assert(rightKeys === Seq(exprA, exprC, exprC)) assert(p.expressions == Seq(exprA, exprC, exprA)) @@ -388,7 +388,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB, exprB)) assert(rightKeys === Seq(exprA, exprC, exprD)) assert(p.expressions == Seq(exprA, exprC, exprA)) @@ -444,8 +444,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == 5) assert(right.numPartitions == 5) case other => fail(other.toString) @@ -462,7 +462,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, DummySparkPlan(_, _, left: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == 10) assert(right.numPartitions == 10) assert(right.expressions == Seq(exprC, exprD)) @@ -481,8 +481,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == 5) assert(left.expressions == Seq(exprA, exprB)) assert(right.numPartitions == 5) @@ -492,7 +492,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, DummySparkPlan(_, _, left: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == 1) assert(right.numPartitions == 1) assert(right.expressions == Seq(exprC)) @@ -510,8 +510,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == conf.numShufflePartitions) assert(left.expressions == Seq(exprA, exprB)) assert(right.numPartitions == conf.numShufflePartitions) @@ -529,7 +529,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, DummySparkPlan(_, _, left: PartitioningCollection, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == 10) assert(right.numPartitions == 10) assert(right.expressions == Seq(exprA)) @@ -545,7 +545,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { Inner, None, plan1, plan2) applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, right: PartitioningCollection, _, _), _), _) => assert(left.numPartitions == 20) assert(left.expressions == Seq(exprC)) @@ -587,7 +587,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, DummySparkPlan(_, _, left: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB)) assert(right.expressions === Seq(exprC, exprD)) assert(left.numPartitions == 6) @@ -606,7 +606,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, right: HashPartitioning, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB)) assert(right.expressions === Seq(exprC, exprD)) @@ -624,7 +624,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, right: HashPartitioning, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB)) assert(right.expressions === Seq(exprC, exprD)) @@ -644,7 +644,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, right: HashPartitioning, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB)) assert(right.expressions === Seq(exprC, exprD)) @@ -666,8 +666,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB)) assert(right.expressions === Seq(exprC, exprD)) assert(left.numPartitions == conf.numShufflePartitions) @@ -691,7 +691,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, left: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB)) assert(rightKeys === Seq(exprC, exprD)) assert(left.numPartitions == 9) @@ -714,8 +714,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { var smjExec = SortMergeJoinExec(exprA :: Nil, exprC :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA)) assert(rightKeys === Seq(exprC)) assert(left.numPartitions == 20) @@ -733,7 +733,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA)) assert(rightKeys === Seq(exprC)) assert(right.numPartitions == 10) @@ -765,8 +765,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { } else { EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA)) assert(rightKeys === Seq(exprC)) assert(left.numPartitions == 5) @@ -875,8 +875,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB, exprB)) assert(right.expressions === Seq(exprA, exprC, exprC)) case other => fail(other.toString) @@ -951,8 +951,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: exprC :: Nil, exprA :: exprB :: exprC :: Nil, Inner, None, plan1, plan2) applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB, exprC)) assert(right.expressions === Seq(exprA, exprB, exprC)) case other => fail(other.toString) @@ -969,8 +969,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB, exprB)) assert(right.expressions === Seq(exprA, exprC, exprC)) case other => fail(other.toString) @@ -987,8 +987,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB, exprB)) assert(right.expressions === Seq(exprA, exprC, exprC)) case other => fail(other.toString) @@ -1008,8 +1008,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB, exprB)) assert(right.expressions === Seq(exprA, exprC, exprC)) case other => fail(other.toString) @@ -1130,7 +1130,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { case ShuffledHashJoinExec(_, _, _, _, _, DummySparkPlan(_, _, left: KeyedPartitioning, _, _), ShuffleExchangeExec(KeyedPartitioning(attrs, pks, _, _), - DummySparkPlan(_, _, SinglePartition, _, _), _, _), _) => + DummySparkPlan(_, _, SinglePartition, _, _), _, _, _), _) => assert(left.expressions == a1 :: Nil) assert(attrs == a1 :: Nil) assert(partitionKeys == pks.map(_.row)) @@ -1234,8 +1234,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(join) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(p1: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p2: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p1: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(p2: HashPartitioning, _, _, _, _), _), _) => // Both sides should be shuffled to default partitions assert(p1.numPartitions == 10) assert(p2.numPartitions == 10) @@ -1259,8 +1259,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(join) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), _) => // Both sides shuffled due to key mismatch case other => fail(s"Expected shuffles on both sides, but got: $other") } @@ -1278,7 +1278,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(join) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), SortExec(_, _, _: DummySparkPlan, _), _) => // Left side shuffled, right side kept as-is case other => fail(s"Expected shuffle on the left side, but got: $other") @@ -1296,8 +1296,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(join) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), _) => // Both sides shuffled due to canCreatePartitioning = false case other => fail(s"Expected shuffles on both sides, but got: $other") } @@ -1371,8 +1371,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(join) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(p1: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p2: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p1: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(p2: HashPartitioning, _, _, _, _), _), _) => // Both sides should be shuffled because partition keys not in join keys assert(p1.numPartitions == 10) assert(p2.numPartitions == 10) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala index c46f0076721b9..c0983f338abe5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala @@ -677,8 +677,8 @@ abstract class StreamingInnerJoinBase extends StreamingJoinSuite { assert(query.lastExecution.executedPlan.collect { case j @ StreamingSymmetricHashJoinExec(_, _, _, _, _, _, _, _, _, - ShuffleExchangeExec(opA: HashPartitioning, _, _, _), - ShuffleExchangeExec(opB: HashPartitioning, _, _, _), _) + ShuffleExchangeExec(opA: HashPartitioning, _, _, _, _), + ShuffleExchangeExec(opB: HashPartitioning, _, _, _, _), _) if partitionExpressionsColumns(opA.expressions) === Seq("a", "b") && partitionExpressionsColumns(opB.expressions) === Seq("a", "b") && opA.numPartitions == numPartitions && opB.numPartitions == numPartitions => j From 6e0c2c26c182b9ed7113d0f534c34a207050260f Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Wed, 22 Jul 2026 04:43:20 +0000 Subject: [PATCH 12/24] [SPARK-XXXXX][SS] Run stateful Real-Time Mode queries over a pipelined shuffle Enable a stateful streaming query (starting with deduplication) to run in Real-Time Mode by pipelining its repartition shuffle: - IncrementalExecution adds a Real-Time Mode preparation rule that marks every streaming shuffle exchange pipelined (ShuffleExchangeExec.pipelined = true). Real-Time Mode is detected structurally by a RealTimeStreamScanExec leaf. The producer (source scan) and consumer (stateful operator) stages are then co-scheduled and stream records through a transient shuffle instead of the consumer waiting for the producer to fully materialize. Inert for a non-RTM batch, so the ordinary microbatch path is unchanged. (AQE is already disabled for RTM batches upstream, so the pipelined shuffle is only ever submitted as a result job, never an AQE map-stage job.) - RealTimeModeAllowlist admits the operators a streaming dropDuplicates plans into (StreamingDeduplicateExec, StateStoreRestoreExec, StateStoreSaveExec) plus ShuffleExchangeExec, so the query is no longer rejected before it runs. Co-authored-by: Isaac --- .../runtime/IncrementalExecution.scala | 32 ++++++++++++++- .../runtime/RealTimeModeAllowlist.scala | 10 +++++ .../StreamRealTimeModeAllowlistSuite.scala | 39 +++---------------- 3 files changed, 45 insertions(+), 36 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala index 0d2e4a6941a00..0f2db3611208d 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala @@ -36,8 +36,9 @@ import org.apache.spark.sql.classic.SparkSession import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.execution.{CommandExecutionMode, LocalLimitExec, QueryExecution, SerializeFromObjectExec, SparkPlan, SparkPlanner, SparkStrategy => Strategy, UnaryExecNode} import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, MergingSessionsExec, ObjectHashAggregateExec, SortAggregateExec, UpdatingSessionsExec} +import org.apache.spark.sql.execution.datasources.v2.RealTimeStreamScanExec import org.apache.spark.sql.execution.datasources.v2.state.metadata.StateMetadataPartitionReader -import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike +import org.apache.spark.sql.execution.exchange.{ShuffleExchangeExec, ShuffleExchangeLike} import org.apache.spark.sql.execution.python.streaming.{FlatMapGroupsInPandasWithStateExec, TransformWithStateInPySparkExec} import org.apache.spark.sql.execution.streaming.{StreamingErrors, StreamingQueryPlanTraverseHelper} import org.apache.spark.sql.execution.streaming.checkpointing.{CheckpointFileManager, OffsetSeqMetadata, OffsetSeqMetadataBase} @@ -663,7 +664,34 @@ class IncrementalExecution( } } - override def preparations: Seq[Rule[SparkPlan]] = state +: super.preparations + /** + * For a Real-Time Mode batch, mark the shuffle exchanges as pipelined so the DAGScheduler + * co-schedules a stateful query's producer (source scan) and consumer (stateful operator) stages + * as one pipelined group -- records stream through a transient shuffle instead of the consumer + * waiting for the producer to fully materialize. The exchange carries the decision as a field + * (see ShuffleExchangeExec.pipelined); the PipelinedShuffleDependency it then builds is the whole + * opt-in -- routing to the streaming shuffle manager and pipelined-group scheduling both follow + * from that dependency type. + * + * Real-Time Mode is detected structurally by a RealTimeStreamScanExec leaf (there is no + * RTM-specific plan flag). Inert for a non-RTM batch, so the ordinary microbatch path is + * unchanged. + */ + object MarkPipelinedShuffleForRealTimeMode extends Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = { + val isRealTimeMode = plan.exists(_.isInstanceOf[RealTimeStreamScanExec]) + if (!isRealTimeMode) { + plan + } else { + plan.transformUp { + case s: ShuffleExchangeExec if !s.pipelined => s.copy(pipelined = true) + } + } + } + } + + override def preparations: Seq[Rule[SparkPlan]] = + state +: (super.preparations :+ MarkPipelinedShuffleForRealTimeMode) /** no need to try-catch again as this is already done once */ override def assertAnalyzed(): Unit = analyzed diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala index 7ae557797f36c..047ea3de4298d 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala @@ -55,7 +55,17 @@ object RealTimeModeAllowlist extends Logging { "org.apache.spark.sql.execution.datasources.v2.WriteToDataSourceV2Exec", "org.apache.spark.sql.execution.exchange.BroadcastExchangeExec", "org.apache.spark.sql.execution.exchange.ReusedExchangeExec", + // A pipelined (streaming) shuffle repartitions a stateful query's input by key. In Real-Time + // Mode the DAGScheduler co-schedules the shuffle's producer and consumer stages via a + // PipelinedShuffleDependency (see IncrementalExecution's pipelined-shuffle rule), so the + // exchange is a supported member of a pipelined group rather than a materialization barrier. + "org.apache.spark.sql.execution.exchange.ShuffleExchangeExec", "org.apache.spark.sql.execution.joins.BroadcastHashJoinExec", + // Streaming deduplication and the state-store access operators it plans into. These run in the + // pipelined-shuffle consumer stage, keyed by the same columns the shuffle repartitions on. + "org.apache.spark.sql.execution.streaming.operators.stateful.StateStoreRestoreExec", + "org.apache.spark.sql.execution.streaming.operators.stateful.StateStoreSaveExec", + "org.apache.spark.sql.execution.streaming.operators.stateful.StreamingDeduplicateExec", classOf[EventTimeWatermarkExec].getName ) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala index a9b8d54a6c61b..bad7cd2cf695f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala @@ -70,7 +70,6 @@ class StreamRealTimeModeAllowlistSuite extends StreamRealTimeModeE2ESuiteBase { "errorType" -> "operator", "message" -> ( "org.apache.spark.sql.execution.SortExec, " + - "org.apache.spark.sql.execution.exchange.ShuffleExchangeExec, " + "org.apache.spark.sql.execution.joins.SortMergeJoinExec are" ) ) @@ -107,31 +106,9 @@ class StreamRealTimeModeAllowlistSuite extends StreamRealTimeModeE2ESuiteBase { } } - // TODO(SPARK-54237) : Remove this test after RTM can shuffle to multiple stages - test("repartition not allowed") { - val inputData = LowLatencyMemoryStream[Int](2) - - val df = inputData.toDF() - .select(col("value").as("key")) - .repartition(4, col("key")) - - val query = runStreamingQuery("repartition_allowlist", df) - - eventually(timeout(60.seconds)) { - checkError( - exception = query.exception.get.getCause.asInstanceOf[SparkIllegalArgumentException], - condition = "STREAMING_REAL_TIME_MODE.OPERATOR_OR_SINK_NOT_IN_ALLOWLIST", - parameters = Map( - "errorType" -> "operator", - "message" -> ( - "org.apache.spark.sql.execution.exchange.ShuffleExchangeExec is" - ) - ) - ) - } - } - - // TODO(SPARK-54236) : Remove this test after RTM supports stateful queries + // A stateful aggregation is still rejected under RTM because HashAggregateExec is not in the + // allowlist, even though the pipelined shuffle and the StateStore operators now are. When RTM + // supports aggregation (SPARK-54236), remove this test. test("stateful queries not allowed") { val inputData = LowLatencyMemoryStream[Int](2) @@ -141,7 +118,7 @@ class StreamRealTimeModeAllowlistSuite extends StreamRealTimeModeE2ESuiteBase { .count() .select(concat(col("key"), lit("-"), col("count"))) - val query = runStreamingQuery("repartition_allowlist", df) + val query = runStreamingQuery("stateful_allowlist", df) eventually(timeout(60.seconds)) { checkError( @@ -149,13 +126,7 @@ class StreamRealTimeModeAllowlistSuite extends StreamRealTimeModeE2ESuiteBase { condition = "STREAMING_REAL_TIME_MODE.OPERATOR_OR_SINK_NOT_IN_ALLOWLIST", parameters = Map( "errorType" -> "operator", - "message" -> ( - "org.apache.spark.sql.execution.aggregate.HashAggregateExec, " + - "org.apache.spark.sql.execution.exchange.ShuffleExchangeExec, " + - "org.apache.spark.sql.execution.streaming" + - ".operators.stateful.StateStoreRestoreExec, " + - "org.apache.spark.sql.execution.streaming.operators.stateful.StateStoreSaveExec are" - ) + "message" -> "org.apache.spark.sql.execution.aggregate.HashAggregateExec is" ) ) } From 2b6616ea7f1732b583777ec0fc47d1e352b5d102 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Wed, 22 Jul 2026 04:50:25 +0000 Subject: [PATCH 13/24] [SPARK-XXXXX][SS][TESTS] Test stateful dedup running in Real-Time Mode over a pipelined shuffle Add an end-to-end test: a streaming dropDuplicates query in Real-Time Mode whose repartition shuffle is pipelined. It asserts correct dedup output across two batches, that every shuffle exchange in the executed plan is pipelined, and that the producer and consumer stages are co-scheduled (>= 2 stages running at once). The Real-Time Mode operator allowlist check is left enabled, so the test also covers admission of the dedup plan's operators. Co-authored-by: Isaac --- .../RealTimeModePipelinedShuffleSuite.scala | 98 +++++++++++++++++++ 1 file changed, 98 insertions(+) create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala new file mode 100644 index 0000000000000..4ddeda78ee037 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala @@ -0,0 +1,98 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.streaming + +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicInteger + +import org.apache.spark.scheduler.{SparkListener, SparkListenerStageCompleted} +import org.apache.spark.scheduler.SparkListenerStageSubmitted +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.sql.execution.streaming.sources.{ContinuousMemorySink, LowLatencyMemoryStream} + +/** + * Tests that a stateful Real-Time Mode query (streaming `dropDuplicates`) runs when its repartition + * shuffle is a PipelinedShuffleDependency, so the source-scan producer stage and the dedup consumer + * stage are co-scheduled and stream records through a transient shuffle instead of the consumer + * waiting for the producer to fully materialize. + * + * The Real-Time Mode operator allowlist check is left enabled, so the test also confirms that the + * dedup plan's operators (StreamingDeduplicateExec, StateStoreRestore/Save, ShuffleExchangeExec) + * are admitted. It drives the source with the standard `testStream` DSL, which advances the query + * across multiple batches. + */ +class RealTimeModePipelinedShuffleSuite extends StreamRealTimeModeSuiteBase { + import testImplicits._ + + test("stateful dedup runs in Real-Time Mode over a pipelined shuffle") { + // Track, from the driver, whether the producer (source scan) and consumer (dedup) stages of the + // pipelined group were ever RUNNING simultaneously. runningStages holds currently-running stage + // ids (added on submit, removed on completion); maxConcurrentStages records the peak. A + // sequential producer-then-consumer schedule never exceeds one running stage at a time. + val runningStages = ConcurrentHashMap.newKeySet[Int]() + val maxConcurrentStages = new AtomicInteger(0) + val listener = new SparkListener { + override def onStageSubmitted(e: SparkListenerStageSubmitted): Unit = { + runningStages.add(e.stageInfo.stageId) + maxConcurrentStages.accumulateAndGet(runningStages.size(), Math.max) + } + override def onStageCompleted(e: SparkListenerStageCompleted): Unit = { + runningStages.remove(e.stageInfo.stageId) + } + } + spark.sparkContext.addSparkListener(listener) + + try { + val inputData = LowLatencyMemoryStream[(String, Int)] + + // scan --shuffle(repartition by key)--> streaming dropDuplicates --> sink. + // dropDuplicates on "key" forces a hash-partitioning ShuffleExchangeExec, which the + // IncrementalExecution rule marks pipelined for a Real-Time Mode batch. + val result = inputData + .toDF() + .select($"_1".as("key")) + .dropDuplicates("key") + .select($"key") + + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + // Batch 0: three distinct keys, each sent twice -> dedup emits each once. + AddData(inputData, ("a", 1), ("b", 1), ("c", 1), ("a", 2), ("b", 2), ("c", 2)), + StartStream(), + CheckAnswer("a", "b", "c"), + // Batch 1: all duplicates of already-seen keys plus one new key -> only "d" is new. + AddData(inputData, ("a", 3), ("b", 3), ("c", 3), ("d", 1)), + CheckAnswer("a", "b", "c", "d"), + // Every Real-Time Mode shuffle exchange in the executed plan is pipelined, and the producer + // + consumer stages were genuinely co-scheduled (>= 2 stages ran at once). + Execute { q => + val executedPlan = q.lastExecution.executedPlan + val exchanges = executedPlan.collect { case s: ShuffleExchangeExec => s } + assert(exchanges.nonEmpty, "expected a shuffle exchange in the dedup plan") + assert(exchanges.forall(_.pipelined), + "expected all Real-Time Mode shuffle exchanges to be pipelined, got: " + + exchanges.map(e => s"pipelined=${e.pipelined}").mkString(", ")) + assert(maxConcurrentStages.get() >= 2, + s"expected >= 2 stages running concurrently, saw max ${maxConcurrentStages.get()}") + }, + StopStream + ) + } finally { + spark.sparkContext.removeSparkListener(listener) + } + } +} From b865ccfee30a3c5f36a3b7f863a404a02de91701 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Wed, 22 Jul 2026 05:08:32 +0000 Subject: [PATCH 14/24] [SPARK-XXXXX][SS][TESTS] Test dedup fault-tolerance in Real-Time Mode over a pipelined shuffle Add a fault-tolerance case to the pipelined-shuffle Real-Time Mode suite: a stateful dropDuplicates query whose repartition shuffle is pipelined, that hits a task failure mid-batch and then restarts from checkpoint. It asserts the query fails on the task error (Real-Time Mode does not retry tasks; a member failure aborts the whole co-scheduled group), and that on restart the deduplication state recovered from the last committed batch -- already-seen keys are not re-emitted. Switches the suite to the manual-clock base so batch boundaries are deterministic across the failure and restart, and uses CheckAnswerWithTimeout (which polls the sink in real time) with advanceRealTimeClock to drive Real-Time Mode batches. Co-authored-by: Isaac --- .../RealTimeModePipelinedShuffleSuite.scala | 78 +++++++++++++++++-- 1 file changed, 72 insertions(+), 6 deletions(-) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala index 4ddeda78ee037..e5fa82863d612 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala @@ -20,10 +20,17 @@ package org.apache.spark.sql.streaming import java.util.concurrent.ConcurrentHashMap import java.util.concurrent.atomic.AtomicInteger +import org.apache.spark.SparkException import org.apache.spark.scheduler.{SparkListener, SparkListenerStageCompleted} import org.apache.spark.scheduler.SparkListenerStageSubmitted import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.execution.streaming.sources.{ContinuousMemorySink, LowLatencyMemoryStream} +import org.apache.spark.sql.functions.udf + +/** Driver-side switch a UDF reads on executors to fail a task on demand (for fault-tolerance). */ +object RealTimeModePipelinedShuffleSuite { + @volatile var failTasks: Boolean = false +} /** * Tests that a stateful Real-Time Mode query (streaming `dropDuplicates`) runs when its repartition @@ -31,14 +38,19 @@ import org.apache.spark.sql.execution.streaming.sources.{ContinuousMemorySink, L * stage are co-scheduled and stream records through a transient shuffle instead of the consumer * waiting for the producer to fully materialize. * - * The Real-Time Mode operator allowlist check is left enabled, so the test also confirms that the + * The Real-Time Mode operator allowlist check is left enabled, so the tests also confirm that the * dedup plan's operators (StreamingDeduplicateExec, StateStoreRestore/Save, ShuffleExchangeExec) - * are admitted. It drives the source with the standard `testStream` DSL, which advances the query + * are admitted. They drive the source with the standard `testStream` DSL, which advances the query * across multiple batches. */ -class RealTimeModePipelinedShuffleSuite extends StreamRealTimeModeSuiteBase { +class RealTimeModePipelinedShuffleSuite extends StreamRealTimeModeManualClockSuiteBase { import testImplicits._ + override def beforeEach(): Unit = { + super.beforeEach() + RealTimeModePipelinedShuffleSuite.failTasks = false + } + test("stateful dedup runs in Real-Time Mode over a pipelined shuffle") { // Track, from the driver, whether the producer (source scan) and consumer (dedup) stages of the // pipelined group were ever RUNNING simultaneously. runningStages holds currently-running stage @@ -70,13 +82,17 @@ class RealTimeModePipelinedShuffleSuite extends StreamRealTimeModeSuiteBase { .select($"key") testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( - // Batch 0: three distinct keys, each sent twice -> dedup emits each once. + // Batch 0: three distinct keys, each sent twice -> dedup emits each once. A Real-Time Mode + // batch runs for a fixed duration and emits in real time, so CheckAnswerWithTimeout polls + // the sink (rather than blocking for batch completion) and advanceRealTimeClock ends it. AddData(inputData, ("a", 1), ("b", 1), ("c", 1), ("a", 2), ("b", 2), ("c", 2)), StartStream(), - CheckAnswer("a", "b", "c"), + CheckAnswerWithTimeout(60000, "a", "b", "c"), + advanceRealTimeClock, + WaitUntilBatchProcessed(0), // Batch 1: all duplicates of already-seen keys plus one new key -> only "d" is new. AddData(inputData, ("a", 3), ("b", 3), ("c", 3), ("d", 1)), - CheckAnswer("a", "b", "c", "d"), + CheckAnswerWithTimeout(60000, "a", "b", "c", "d"), // Every Real-Time Mode shuffle exchange in the executed plan is pipelined, and the producer // + consumer stages were genuinely co-scheduled (>= 2 stages ran at once). Execute { q => @@ -95,4 +111,54 @@ class RealTimeModePipelinedShuffleSuite extends StreamRealTimeModeSuiteBase { spark.sparkContext.removeSparkListener(listener) } } + + test("dedup over a pipelined shuffle recovers from a task failure via checkpoint restart") { + withTempDir { checkpointDir => + // A UDF that throws on demand, to fail a task mid-batch. Placed after the dedup so the + // failure lands in the pipelined consumer stage while the query is running. + val failUDF = udf { (key: String) => + if (RealTimeModePipelinedShuffleSuite.failTasks) { + throw new RuntimeException(s"forced task failure on $key") + } + key + } + + val inputData = LowLatencyMemoryStream[(String, Int)] + val result = inputData + .toDF() + .select($"_1".as("key")) + .dropDuplicates("key") + .select(failUDF($"key").as("key")) + + // First run: dedup "a","b" (batch 0 commits), then a task failure fails the query. RTM does + // not retry tasks, so a single task failure fails the whole batch/query. + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, ("a", 1), ("b", 1), ("a", 2)), + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + CheckAnswerWithTimeout(60000, "a", "b"), + advanceRealTimeClock, + WaitUntilBatchProcessed(0), + Execute { _ => RealTimeModePipelinedShuffleSuite.failTasks = true }, + AddData(inputData, ("c", 1)), + advanceRealTimeClock, + ExpectFailure[SparkException] { ex => + val msg = Option(ex.getCause).map(_.getMessage).getOrElse(ex.getMessage) + assert(msg != null && msg.contains("forced task failure"), + s"expected a forced task failure, got: $msg") + } + ) + + // Restart from the same checkpoint. The dedup state from batch 0 must survive: "a" and "b" + // are already seen and must NOT be re-emitted; only the genuinely-new "c","d" appear. This + // proves recovery-to-last-committed-batch works when the shuffle is pipelined. + RealTimeModePipelinedShuffleSuite.failTasks = false + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + AddData(inputData, ("a", 3), ("b", 3), ("c", 3), ("d", 1)), + CheckAnswerWithTimeout(60000, "c", "d"), + advanceRealTimeClock, + StopStream + ) + } + } } From 5163533a6d7166de2a06ed2fbae13bb9c4248638 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Wed, 22 Jul 2026 06:33:13 +0000 Subject: [PATCH 15/24] [SPARK-XXXXX][SS][TESTS] Test dedup checkpoint-write-failure recovery in Real-Time Mode Add a second fault-tolerance case: a dedup query over a pipelined shuffle whose commit log write is fault-injected (via FailureInjectionCheckpointFileManager) so a batch cannot commit and the query fails. On restart from the same checkpoint, the committed batch's deduplication state survives and the uncommitted batch is re-run, so its data is still emitted. The injected close() failure surfaces as the underlying IOException in this path (the COMMIT_LOG_WRITE_FAILURE error-class wrapping is not open source), so the test expects that. Co-authored-by: Isaac --- .../RealTimeModePipelinedShuffleSuite.scala | 59 +++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala index e5fa82863d612..cf6abeff4884e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala @@ -17,6 +17,7 @@ package org.apache.spark.sql.streaming +import java.io.{File, IOException} import java.util.concurrent.ConcurrentHashMap import java.util.concurrent.atomic.AtomicInteger @@ -25,7 +26,9 @@ import org.apache.spark.scheduler.{SparkListener, SparkListenerStageCompleted} import org.apache.spark.scheduler.SparkListenerStageSubmitted import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.execution.streaming.sources.{ContinuousMemorySink, LowLatencyMemoryStream} +import org.apache.spark.sql.execution.streaming.state.{FailureInjectionCheckpointFileManager, FailureInjectionFileSystem, FailureInjectionState} import org.apache.spark.sql.functions.udf +import org.apache.spark.sql.internal.SQLConf /** Driver-side switch a UDF reads on executors to fail a task on demand (for fault-tolerance). */ object RealTimeModePipelinedShuffleSuite { @@ -161,4 +164,60 @@ class RealTimeModePipelinedShuffleSuite extends StreamRealTimeModeManualClockSui ) } } + + /** Run `f` with a temp dir whose checkpoint file ops can be fault-injected via injectionState. */ + private def withTempDirAllowFailureInjection(f: (File, FailureInjectionState) => Unit): Unit = { + withTempDir { dir => + val injectionState = FailureInjectionFileSystem.registerTempPath(dir.getPath) + try { + f(dir, injectionState) + } finally { + FailureInjectionFileSystem.removePathFromTempToInjectionState(dir.getPath) + } + } + } + + test("dedup over a pipelined shuffle recovers from a commit-log write failure") { + withSQLConf( + SQLConf.STREAMING_CHECKPOINT_FILE_MANAGER_CLASS.parent.key -> + classOf[FailureInjectionCheckpointFileManager].getName) { + withTempDirAllowFailureInjection { (checkpointDir, injectionState) => + val inputData = LowLatencyMemoryStream[(String, Int)] + val result = inputData + .toDF() + .select($"_1".as("key")) + .dropDuplicates("key") + .select($"key") + + // Batch 0 dedups "a","b" and commits. Then fail the close() of batch 1's commit-log write, + // so batch 1 cannot commit and the query fails after processing. + injectionState.createAtomicDelayCloseRegex = Seq(".*/commits/1") + + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, ("a", 1), ("b", 1)), + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + CheckAnswerWithTimeout(60000, "a", "b"), + advanceRealTimeClock, + WaitUntilBatchProcessed(0), + AddData(inputData, ("c", 1)), + CheckAnswerWithTimeout(60000, "a", "b", "c"), + advanceRealTimeClock, + // The injected close() failure surfaces as the underlying IOException failing the batch. + ExpectFailure[IOException]() + ) + + // Clear the injection and restart from the same checkpoint. Batch 0's dedup state must + // survive (a, b already seen); the uncommitted batch 1 is re-run, so its new key "c" is + // still emitted, and a further new key "d" is added. + injectionState.createAtomicDelayCloseRegex = Seq.empty + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + AddData(inputData, ("a", 2), ("b", 2), ("c", 2), ("d", 1)), + CheckAnswerWithTimeout(60000, "c", "d"), + advanceRealTimeClock, + StopStream + ) + } + } + } } From 9fc2bd88f6cea50b513091d10b50b5f202166be8 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Wed, 29 Jul 2026 19:09:57 +0000 Subject: [PATCH 16/24] [SPARK-XXXXX][SS] Document that RTM marks every shuffle as pipelined MarkPipelinedShuffleForRealTimeMode marks every shuffle exchange in an RTM plan as pipelined, so a chain of several pipelined shuffles (e.g. two repartitions, or a repartition feeding a keyed stateful operator) is co-scheduled as one pipelined group. The one unsupported shape -- a pipelined shuffle with more than one consumer (fan-out, e.g. a reused exchange) -- is rejected up front by the DAGScheduler, so no ReusedExchangeExec reasoning is needed here. Co-authored-by: Isaac --- .../streaming/runtime/IncrementalExecution.scala | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala index 0f2db3611208d..18644e0e780c2 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala @@ -676,6 +676,22 @@ class IncrementalExecution( * Real-Time Mode is detected structurally by a RealTimeStreamScanExec leaf (there is no * RTM-specific plan flag). Inert for a non-RTM batch, so the ordinary microbatch path is * unchanged. + * + * Marks EVERY shuffle exchange, so a plan with several pipelined shuffles in a chain (e.g. two + * repartitions, or a repartition feeding a keyed stateful operator) is handled: each becomes a + * PipelinedShuffleDependency and the whole all-pipelined job is co-scheduled as one pipelined + * group (the DAGScheduler treats an all-pipelined job's stage graph as a single group). No + * shuffle-count restriction is needed here. + * + * transformUp does not descend into a ReusedExchangeExec (a leaf whose wrapped exchange is a + * field, not a tree child), so a REUSED shuffle exchange would keep pipelined=false while its + * standalone twin flips to true. That divergence is not reachable: a reused shuffle requires + * referencing the same streaming source more than once (self-join / self-union / CTE read twice), + * which Real-Time Mode rejects up front (MicroBatchExecution, IDENTICAL_SOURCES_IN_UNION_NOT_ + * SUPPORTED) before this rule runs. The only ReusedExchangeExec that reaches an RTM plan wraps a + * BROADCAST exchange (multiple broadcast joins on the same static table, SC-209926), which this + * rule does not match. A pipelined shuffle read by more than one consumer (fan-out) is likewise + * rejected up front by the DAGScheduler (checkPipelinedGroupsSupportedInRDDGraph). */ object MarkPipelinedShuffleForRealTimeMode extends Rule[SparkPlan] { override def apply(plan: SparkPlan): SparkPlan = { From 7a98316b8894f1621c754273504837ccf64a529e Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sun, 2 Aug 2026 00:07:28 +0000 Subject: [PATCH 17/24] [SPARK-XXXXX][SS] Disable sort before repartition in Real-Time Mode `repartition(n)` uses RoundRobinPartitioning, and by default (SPARK-23207) Spark inserts a local sort before a round-robin shuffle so the row-to-partition assignment is deterministic across task retries. That sort is fully blocking: it drains its entire input before emitting a row. A Real-Time Mode task reads an unbounded stream, so the input never ends. The producer stage therefore never emits a row, never reaches the shuffle write, and every downstream stage waits forever -- the query returns no rows at all. This is not visible to the Real-Time Mode operator allowlist because the sort is not a SortExec node; it lives inside ShuffleExchangeExec's RDD construction. Determinism is not needed here: a pipelined task set is capped at a single attempt (TaskSetManager), and a pipelined group is aborted as a unit rather than recomputed, so the retry hazard the sort guards against cannot arise. Disable the config for Real-Time Mode unless the user has set it explicitly. Co-authored-by: Isaac --- .../streaming/runtime/MicroBatchExecution.scala | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) 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 de84c9d15a6e4..f608abd4423a6 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 @@ -371,6 +371,23 @@ class MicroBatchExecution( messageParameters = Map("className" -> s.getClass.getName) ) } + + // `repartition(n)` uses RoundRobinPartitioning, and by default (SPARK-23207) Spark inserts a + // local sort before a round-robin shuffle so the row-to-partition assignment is deterministic + // -- otherwise a retried task could emit rows in a different order and lose data. That sort + // is fully blocking: it drains its whole input before emitting a row. A Real-Time Mode task + // reads an unbounded stream, so the input never ends, the producer never emits, and the query + // makes no progress at all. Note this sort is not a SortExec node (it lives inside + // ShuffleExchangeExec's RDD), so the Real-Time Mode operator allowlist cannot catch it. + // + // Determinism is not needed here: Real-Time Mode does not retry tasks (TaskSetManager caps a + // pipelined task set at one attempt, and its group is aborted as a unit rather than + // recomputed), so the hazard the sort guards against does not arise. Only set the config when + // the user has not made an explicit choice. + if (!sparkSessionForStream.conf.contains(SQLConf.SORT_BEFORE_REPARTITION.key)) { + logInfo("Disabling sort before repartition for Real-Time Mode") + sparkSessionForStream.conf.set(SQLConf.SORT_BEFORE_REPARTITION.key, "false") + } } // Initializing TriggerExecutor relies on `sources`, hence calling this after initializing From c9b57667641990fcf3eb77ca2ece15838a509c2d Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sun, 2 Aug 2026 00:07:43 +0000 Subject: [PATCH 18/24] [SPARK-XXXXX][SS] Move RTM pipelined-shuffle tests into the shared suite and add coverage Move the pipelined-shuffle Real-Time Mode tests out of the standalone RealTimeModePipelinedShuffleSuite into StreamRealTimeModeSuite.scala, which already exists in both this repo and the Databricks runtime. The clock-driven tests land in StreamRealTimeModeWithManualClockSuite, whose base supplies advanceRealTimeClock; the reused-broadcast test needs no manual clock and goes in the parent suite. This keeps the tests in a file that both sides share. New coverage: - A chain of two pipelined shuffles (a round-robin repartition feeding a keyed dropDuplicates) forms a single co-scheduled pipelined group. This is the multi-shuffle-per-group shape, previously untested, and is the regression test for the blocking sort fixed in the previous commit. - Multi-key dropDuplicates over a pipelined shuffle. - An explicit repartition-by-key in Real-Time Mode. - Multiple broadcast joins against the same static table, which produces a ReusedExchangeExec (SC-209926). Co-authored-by: Isaac --- .../RealTimeModePipelinedShuffleSuite.scala | 223 --------------- .../streaming/StreamRealTimeModeSuite.scala | 264 +++++++++++++++++- 2 files changed, 260 insertions(+), 227 deletions(-) delete mode 100644 sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala deleted file mode 100644 index cf6abeff4884e..0000000000000 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeModePipelinedShuffleSuite.scala +++ /dev/null @@ -1,223 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.apache.spark.sql.streaming - -import java.io.{File, IOException} -import java.util.concurrent.ConcurrentHashMap -import java.util.concurrent.atomic.AtomicInteger - -import org.apache.spark.SparkException -import org.apache.spark.scheduler.{SparkListener, SparkListenerStageCompleted} -import org.apache.spark.scheduler.SparkListenerStageSubmitted -import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec -import org.apache.spark.sql.execution.streaming.sources.{ContinuousMemorySink, LowLatencyMemoryStream} -import org.apache.spark.sql.execution.streaming.state.{FailureInjectionCheckpointFileManager, FailureInjectionFileSystem, FailureInjectionState} -import org.apache.spark.sql.functions.udf -import org.apache.spark.sql.internal.SQLConf - -/** Driver-side switch a UDF reads on executors to fail a task on demand (for fault-tolerance). */ -object RealTimeModePipelinedShuffleSuite { - @volatile var failTasks: Boolean = false -} - -/** - * Tests that a stateful Real-Time Mode query (streaming `dropDuplicates`) runs when its repartition - * shuffle is a PipelinedShuffleDependency, so the source-scan producer stage and the dedup consumer - * stage are co-scheduled and stream records through a transient shuffle instead of the consumer - * waiting for the producer to fully materialize. - * - * The Real-Time Mode operator allowlist check is left enabled, so the tests also confirm that the - * dedup plan's operators (StreamingDeduplicateExec, StateStoreRestore/Save, ShuffleExchangeExec) - * are admitted. They drive the source with the standard `testStream` DSL, which advances the query - * across multiple batches. - */ -class RealTimeModePipelinedShuffleSuite extends StreamRealTimeModeManualClockSuiteBase { - import testImplicits._ - - override def beforeEach(): Unit = { - super.beforeEach() - RealTimeModePipelinedShuffleSuite.failTasks = false - } - - test("stateful dedup runs in Real-Time Mode over a pipelined shuffle") { - // Track, from the driver, whether the producer (source scan) and consumer (dedup) stages of the - // pipelined group were ever RUNNING simultaneously. runningStages holds currently-running stage - // ids (added on submit, removed on completion); maxConcurrentStages records the peak. A - // sequential producer-then-consumer schedule never exceeds one running stage at a time. - val runningStages = ConcurrentHashMap.newKeySet[Int]() - val maxConcurrentStages = new AtomicInteger(0) - val listener = new SparkListener { - override def onStageSubmitted(e: SparkListenerStageSubmitted): Unit = { - runningStages.add(e.stageInfo.stageId) - maxConcurrentStages.accumulateAndGet(runningStages.size(), Math.max) - } - override def onStageCompleted(e: SparkListenerStageCompleted): Unit = { - runningStages.remove(e.stageInfo.stageId) - } - } - spark.sparkContext.addSparkListener(listener) - - try { - val inputData = LowLatencyMemoryStream[(String, Int)] - - // scan --shuffle(repartition by key)--> streaming dropDuplicates --> sink. - // dropDuplicates on "key" forces a hash-partitioning ShuffleExchangeExec, which the - // IncrementalExecution rule marks pipelined for a Real-Time Mode batch. - val result = inputData - .toDF() - .select($"_1".as("key")) - .dropDuplicates("key") - .select($"key") - - testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( - // Batch 0: three distinct keys, each sent twice -> dedup emits each once. A Real-Time Mode - // batch runs for a fixed duration and emits in real time, so CheckAnswerWithTimeout polls - // the sink (rather than blocking for batch completion) and advanceRealTimeClock ends it. - AddData(inputData, ("a", 1), ("b", 1), ("c", 1), ("a", 2), ("b", 2), ("c", 2)), - StartStream(), - CheckAnswerWithTimeout(60000, "a", "b", "c"), - advanceRealTimeClock, - WaitUntilBatchProcessed(0), - // Batch 1: all duplicates of already-seen keys plus one new key -> only "d" is new. - AddData(inputData, ("a", 3), ("b", 3), ("c", 3), ("d", 1)), - CheckAnswerWithTimeout(60000, "a", "b", "c", "d"), - // Every Real-Time Mode shuffle exchange in the executed plan is pipelined, and the producer - // + consumer stages were genuinely co-scheduled (>= 2 stages ran at once). - Execute { q => - val executedPlan = q.lastExecution.executedPlan - val exchanges = executedPlan.collect { case s: ShuffleExchangeExec => s } - assert(exchanges.nonEmpty, "expected a shuffle exchange in the dedup plan") - assert(exchanges.forall(_.pipelined), - "expected all Real-Time Mode shuffle exchanges to be pipelined, got: " + - exchanges.map(e => s"pipelined=${e.pipelined}").mkString(", ")) - assert(maxConcurrentStages.get() >= 2, - s"expected >= 2 stages running concurrently, saw max ${maxConcurrentStages.get()}") - }, - StopStream - ) - } finally { - spark.sparkContext.removeSparkListener(listener) - } - } - - test("dedup over a pipelined shuffle recovers from a task failure via checkpoint restart") { - withTempDir { checkpointDir => - // A UDF that throws on demand, to fail a task mid-batch. Placed after the dedup so the - // failure lands in the pipelined consumer stage while the query is running. - val failUDF = udf { (key: String) => - if (RealTimeModePipelinedShuffleSuite.failTasks) { - throw new RuntimeException(s"forced task failure on $key") - } - key - } - - val inputData = LowLatencyMemoryStream[(String, Int)] - val result = inputData - .toDF() - .select($"_1".as("key")) - .dropDuplicates("key") - .select(failUDF($"key").as("key")) - - // First run: dedup "a","b" (batch 0 commits), then a task failure fails the query. RTM does - // not retry tasks, so a single task failure fails the whole batch/query. - testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( - AddData(inputData, ("a", 1), ("b", 1), ("a", 2)), - StartStream(checkpointLocation = checkpointDir.getAbsolutePath), - CheckAnswerWithTimeout(60000, "a", "b"), - advanceRealTimeClock, - WaitUntilBatchProcessed(0), - Execute { _ => RealTimeModePipelinedShuffleSuite.failTasks = true }, - AddData(inputData, ("c", 1)), - advanceRealTimeClock, - ExpectFailure[SparkException] { ex => - val msg = Option(ex.getCause).map(_.getMessage).getOrElse(ex.getMessage) - assert(msg != null && msg.contains("forced task failure"), - s"expected a forced task failure, got: $msg") - } - ) - - // Restart from the same checkpoint. The dedup state from batch 0 must survive: "a" and "b" - // are already seen and must NOT be re-emitted; only the genuinely-new "c","d" appear. This - // proves recovery-to-last-committed-batch works when the shuffle is pipelined. - RealTimeModePipelinedShuffleSuite.failTasks = false - testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( - StartStream(checkpointLocation = checkpointDir.getAbsolutePath), - AddData(inputData, ("a", 3), ("b", 3), ("c", 3), ("d", 1)), - CheckAnswerWithTimeout(60000, "c", "d"), - advanceRealTimeClock, - StopStream - ) - } - } - - /** Run `f` with a temp dir whose checkpoint file ops can be fault-injected via injectionState. */ - private def withTempDirAllowFailureInjection(f: (File, FailureInjectionState) => Unit): Unit = { - withTempDir { dir => - val injectionState = FailureInjectionFileSystem.registerTempPath(dir.getPath) - try { - f(dir, injectionState) - } finally { - FailureInjectionFileSystem.removePathFromTempToInjectionState(dir.getPath) - } - } - } - - test("dedup over a pipelined shuffle recovers from a commit-log write failure") { - withSQLConf( - SQLConf.STREAMING_CHECKPOINT_FILE_MANAGER_CLASS.parent.key -> - classOf[FailureInjectionCheckpointFileManager].getName) { - withTempDirAllowFailureInjection { (checkpointDir, injectionState) => - val inputData = LowLatencyMemoryStream[(String, Int)] - val result = inputData - .toDF() - .select($"_1".as("key")) - .dropDuplicates("key") - .select($"key") - - // Batch 0 dedups "a","b" and commits. Then fail the close() of batch 1's commit-log write, - // so batch 1 cannot commit and the query fails after processing. - injectionState.createAtomicDelayCloseRegex = Seq(".*/commits/1") - - testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( - AddData(inputData, ("a", 1), ("b", 1)), - StartStream(checkpointLocation = checkpointDir.getAbsolutePath), - CheckAnswerWithTimeout(60000, "a", "b"), - advanceRealTimeClock, - WaitUntilBatchProcessed(0), - AddData(inputData, ("c", 1)), - CheckAnswerWithTimeout(60000, "a", "b", "c"), - advanceRealTimeClock, - // The injected close() failure surfaces as the underlying IOException failing the batch. - ExpectFailure[IOException]() - ) - - // Clear the injection and restart from the same checkpoint. Batch 0's dedup state must - // survive (a, b already seen); the uncommitted batch 1 is re-run, so its new key "c" is - // still emitted, and a further new key "d" is added. - injectionState.createAtomicDelayCloseRegex = Seq.empty - testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( - StartStream(checkpointLocation = checkpointDir.getAbsolutePath), - AddData(inputData, ("a", 2), ("b", 2), ("c", 2), ("d", 1)), - CheckAnswerWithTimeout(60000, "c", "d"), - advanceRealTimeClock, - StopStream - ) - } - } - } -} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala index 0f584d818c06e..0bda77e8415dc 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala @@ -17,18 +17,22 @@ package org.apache.spark.sql.streaming -import java.util.concurrent.TimeUnit +import java.io.IOException +import java.util.concurrent.{ConcurrentHashMap, TimeUnit} import java.util.concurrent.atomic.AtomicInteger import scala.concurrent.duration.Duration import org.scalatest.concurrent.PatienceConfiguration.Timeout -import org.apache.spark.{SparkIllegalArgumentException, SparkIllegalStateException, TaskContext} +import org.apache.spark.{SparkException, SparkIllegalArgumentException, SparkIllegalStateException, TaskContext} +import org.apache.spark.scheduler.{SparkListener, SparkListenerStageCompleted, SparkListenerStageSubmitted} +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.execution.streaming.RealTimeTrigger -import org.apache.spark.sql.execution.streaming.runtime.MemoryStream +import org.apache.spark.sql.execution.streaming.runtime.{MemoryStream, StreamExecution} import org.apache.spark.sql.execution.streaming.sources.{ContinuousMemorySink, LowLatencyMemoryStream} -import org.apache.spark.sql.functions.udf +import org.apache.spark.sql.execution.streaming.state.{FailureInjectionCheckpointFileManager, FailureInjectionFileSystem} +import org.apache.spark.sql.functions.{broadcast, concat, lit, udf} import org.apache.spark.sql.internal.SQLConf class StreamRealTimeModeSuite extends StreamRealTimeModeSuiteBase { @@ -209,6 +213,22 @@ class StreamRealTimeModeSuite extends StreamRealTimeModeSuiteBase { StopStream ) } + + test("multiple broadcast joins with the same static table run in Real-Time Mode") { + // Two broadcast joins against the SAME static table -- exchange reuse collapses the second + // broadcast into a ReusedExchangeExec. This is the supported reuse shape in Real-Time Mode: a + // reused BROADCAST exchange (allowlisted), NOT a reused shuffle. The RTM marking rule only + // marks ShuffleExchangeExec, so it correctly leaves the broadcast reuse alone. + val inputData = LowLatencyMemoryStream.singlePartition[Int] + val staticData = Seq((1, "a"), (2, "b"), (3, "c")).toDF("key", "value") + val df = inputData.toDS().toDF("key").join(broadcast(staticData), Seq("key"), "left") + val df2 = df.join(broadcast(staticData), Seq("key"), "left") + testStream(df2, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, 1, 2, 3), + StartStream(), + CheckAnswerWithTimeout(30000, (1, "a", "a"), (2, "b", "b"), (3, "c", "c")) + ) + } } class StreamRealTimeModeWithManualClockSuite extends StreamRealTimeModeManualClockSuiteBase { @@ -393,4 +413,240 @@ class StreamRealTimeModeWithManualClockSuite extends StreamRealTimeModeManualClo StopStream ) } + + // ======================================================================================== + // Pipelined (streaming) shuffle: a stateful/repartition Real-Time Mode query whose shuffle is a + // PipelinedShuffleDependency, so the producer (source scan) and consumer stages are co-scheduled + // and stream records through a transient shuffle instead of the consumer waiting for the producer + // to fully materialize. + // ======================================================================================== + + override def beforeEach(): Unit = { + super.beforeEach() + StreamRealTimeModeSuite.failTasks = false + } + + /** Assert every shuffle exchange in the query's last executed plan is pipelined. */ + private def assertAllExchangesPipelined(q: StreamExecution): Unit = { + val exchanges = q.lastExecution.executedPlan.collect { case s: ShuffleExchangeExec => s } + assert(exchanges.nonEmpty, "expected at least one shuffle exchange in the plan") + assert(exchanges.forall(_.pipelined), + "expected all Real-Time Mode shuffle exchanges to be pipelined, got: " + + exchanges.map(e => s"pipelined=${e.pipelined}").mkString(", ")) + } + + test("pipelined shuffle: stateful dedup runs in Real-Time Mode and co-schedules its stages") { + // Track, from the driver, whether the producer (source scan) and consumer (dedup) stages of the + // pipelined group were ever RUNNING simultaneously. A sequential producer-then-consumer + // schedule never exceeds one running stage at a time; >= 2 proves genuine co-scheduling. + val runningStages = ConcurrentHashMap.newKeySet[Int]() + val maxConcurrentStages = new AtomicInteger(0) + val listener = new SparkListener { + override def onStageSubmitted(e: SparkListenerStageSubmitted): Unit = { + runningStages.add(e.stageInfo.stageId) + maxConcurrentStages.accumulateAndGet(runningStages.size(), Math.max) + } + override def onStageCompleted(e: SparkListenerStageCompleted): Unit = { + runningStages.remove(e.stageInfo.stageId) + } + } + spark.sparkContext.addSparkListener(listener) + try { + val inputData = LowLatencyMemoryStream[(String, Int)] + // scan --shuffle(repartition by key)--> streaming dropDuplicates --> sink. + val result = inputData.toDF().select($"_1".as("key")).dropDuplicates("key").select($"key") + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, ("a", 1), ("b", 1), ("c", 1), ("a", 2), ("b", 2), ("c", 2)), + StartStream(), + CheckAnswerWithTimeout(60000, "a", "b", "c"), + advanceRealTimeClock, + WaitUntilBatchProcessed(0), + AddData(inputData, ("a", 3), ("b", 3), ("c", 3), ("d", 1)), + CheckAnswerWithTimeout(60000, "a", "b", "c", "d"), + Execute { q => + assertAllExchangesPipelined(q) + assert(maxConcurrentStages.get() >= 2, + s"expected >= 2 stages running concurrently, saw max ${maxConcurrentStages.get()}") + }, + StopStream + ) + } finally { + spark.sparkContext.removeSparkListener(listener) + } + } + + test("pipelined shuffle: multi-key dedup runs in Real-Time Mode over a pipelined shuffle") { + // dropDuplicates on more than one column: the hash-partitioning is by the composite key. + // Confirms the pipelined path is not specific to a single dedup column. + val inputData = LowLatencyMemoryStream[(String, Int)] + val result = inputData.toDF().select($"_1".as("k1"), $"_2".as("k2")) + .dropDuplicates("k1", "k2") + .select(concat($"k1", lit("-"), $"k2").as("out")) + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + // (a,1) twice -> once; (a,2) is a distinct composite key -> also emitted. + AddData(inputData, ("a", 1), ("a", 1), ("a", 2), ("b", 1)), + StartStream(), + CheckAnswerWithTimeout(60000, "a-1", "a-2", "b-1"), + advanceRealTimeClock, + WaitUntilBatchProcessed(0), + AddData(inputData, ("a", 1), ("a", 2), ("b", 2)), + CheckAnswerWithTimeout(60000, "a-1", "a-2", "b-1", "b-2"), + Execute { q => assertAllExchangesPipelined(q) }, + StopStream + ) + } + + test("pipelined shuffle: an explicit repartition-by-key runs in Real-Time Mode") { + // A stateless repartition (no dedup): the ShuffleExchangeExec is still marked pipelined in RTM. + val inputData = LowLatencyMemoryStream[Int] + val result = inputData.toDF().repartition(4, $"value").select($"value") + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, 1, 2, 3, 4, 5), + StartStream(), + CheckAnswerWithTimeout(60000, 1, 2, 3, 4, 5), + Execute { q => assertAllExchangesPipelined(q) }, + advanceRealTimeClock, + WaitUntilBatchProcessed(0), + AddData(inputData, 6, 7), + CheckAnswerWithTimeout(60000, 1, 2, 3, 4, 5, 6, 7), + StopStream + ) + } + + test("pipelined shuffle: a chain of two shuffles is a single co-scheduled pipelined group") { + // A round-robin repartition (RoundRobinPartitioning) feeding a dropDuplicates (HashPartitioning + // on the dedup key) produces TWO distinct shuffle exchanges -- neither distribution satisfies + // the other, so they are not collapsed. BOTH must be marked pipelined and the all-pipelined job + // co-schedules as one group (>= 2 stages running at once). This is the multi-shuffle-per- + // group scenario: more than one pipelined shuffle in a single Real-Time Mode job. + val runningStages = ConcurrentHashMap.newKeySet[Int]() + val maxConcurrentStages = new AtomicInteger(0) + val listener = new SparkListener { + override def onStageSubmitted(e: SparkListenerStageSubmitted): Unit = { + runningStages.add(e.stageInfo.stageId) + maxConcurrentStages.accumulateAndGet(runningStages.size(), Math.max) + } + override def onStageCompleted(e: SparkListenerStageCompleted): Unit = { + runningStages.remove(e.stageInfo.stageId) + } + } + spark.sparkContext.addSparkListener(listener) + // Keep the shuffle partition count small: the whole group (scan + 2 shuffles) must fit in the + // test cluster's slots at once (gang admission), so a 2-shuffle chain at the default 200 + // partitions would fail with CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT. + try { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val inputData = LowLatencyMemoryStream[(String, Int)] + val result = inputData.toDF().select($"_1".as("key")) + .repartition(4) // RoundRobinPartitioning -> shuffle #1 + .dropDuplicates("key") // HashPartitioning(key) -> shuffle #2 + .select($"key") + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, ("a", 1), ("b", 1), ("c", 1), ("a", 2)), + StartStream(), + CheckAnswerWithTimeout(60000, "a", "b", "c"), + Execute { q => + val exchanges = q.lastExecution.executedPlan.collect { + case s: ShuffleExchangeExec => s + } + assert(exchanges.size >= 2, s"expected >= 2 shuffle exchanges, got ${exchanges.size}") + assert(exchanges.forall(_.pipelined), + "every exchange in the chain must be pipelined, got: " + + exchanges.map(_.pipelined).mkString(", ")) + assert(maxConcurrentStages.get() >= 2, + s"a 2-shuffle chain must co-schedule >= 2 stages, saw ${maxConcurrentStages.get()}") + }, + StopStream + ) + } + } finally { + spark.sparkContext.removeSparkListener(listener) + } + } + + test("pipelined shuffle: dedup recovers from a task failure via checkpoint restart") { + withTempDir { checkpointDir => + // A UDF that throws on demand, placed after the dedup so the failure lands in the pipelined + // consumer stage while the query runs. RTM does not retry tasks, so one failure fails it. + val failUDF = udf { (key: String) => + if (StreamRealTimeModeSuite.failTasks) { + throw new RuntimeException(s"forced task failure on $key") + } + key + } + val inputData = LowLatencyMemoryStream[(String, Int)] + val result = inputData.toDF().select($"_1".as("key")).dropDuplicates("key") + .select(failUDF($"key").as("key")) + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, ("a", 1), ("b", 1), ("a", 2)), + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + CheckAnswerWithTimeout(60000, "a", "b"), + advanceRealTimeClock, + WaitUntilBatchProcessed(0), + Execute { _ => StreamRealTimeModeSuite.failTasks = true }, + AddData(inputData, ("c", 1)), + advanceRealTimeClock, + ExpectFailure[SparkException] { ex => + val msg = Option(ex.getCause).map(_.getMessage).getOrElse(ex.getMessage) + assert(msg != null && msg.contains("forced task failure"), + s"expected a forced task failure, got: $msg") + } + ) + // Restart from the same checkpoint: batch-0 dedup state must survive (a, b not re-emitted), + // only the genuinely-new c, d appear -- recovery to the last committed batch. + StreamRealTimeModeSuite.failTasks = false + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + AddData(inputData, ("a", 3), ("b", 3), ("c", 3), ("d", 1)), + CheckAnswerWithTimeout(60000, "c", "d"), + advanceRealTimeClock, + StopStream + ) + } + } + + test("pipelined shuffle: dedup recovers from a commit-log write failure") { + withSQLConf( + SQLConf.STREAMING_CHECKPOINT_FILE_MANAGER_CLASS.parent.key -> + classOf[FailureInjectionCheckpointFileManager].getName) { + withTempDir { checkpointDir => + val injectionState = FailureInjectionFileSystem.registerTempPath(checkpointDir.getPath) + try { + val inputData = LowLatencyMemoryStream[(String, Int)] + val result = inputData.toDF().select($"_1".as("key")).dropDuplicates("key").select($"key") + // Batch 0 dedups a, b and commits. Then fail the close() of batch 1's commit-log write so + // batch 1 cannot commit and the query fails after processing. + injectionState.createAtomicDelayCloseRegex = Seq(".*/commits/1") + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, ("a", 1), ("b", 1)), + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + CheckAnswerWithTimeout(60000, "a", "b"), + advanceRealTimeClock, + WaitUntilBatchProcessed(0), + AddData(inputData, ("c", 1)), + CheckAnswerWithTimeout(60000, "a", "b", "c"), + advanceRealTimeClock, + ExpectFailure[IOException]() + ) + // Clear injection, restart: batch-0 state survives (a, b seen); uncommitted batch 1 + // re-runs so its new key c is still emitted, plus a further new key d. + injectionState.createAtomicDelayCloseRegex = Seq.empty + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + AddData(inputData, ("a", 2), ("b", 2), ("c", 2), ("d", 1)), + CheckAnswerWithTimeout(60000, "c", "d"), + advanceRealTimeClock, + StopStream + ) + } finally { + FailureInjectionFileSystem.removePathFromTempToInjectionState(checkpointDir.getPath) + } + } + } + } +} + +/** Driver-side switch a UDF reads on executors to fail a task on demand (fault-tolerance tests). */ +object StreamRealTimeModeSuite { + @volatile var failTasks: Boolean = false } From 037bc865bd4312d5f30711645a170cd4f7ec4a22 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sun, 2 Aug 2026 00:20:58 +0000 Subject: [PATCH 19/24] [SPARK-XXXXX][SS] Do not mark a pipelined round-robin shuffle order-sensitive Turning off the deterministic local sort for Real-Time Mode (previous commit) has a second effect: it flips the same condition's other branch on, marking the map RDD of a round-robin shuffle order-sensitive. That marking exists only so a retry is handled correctly -- an order-sensitive RDD whose own input is UNORDERED becomes INDETERMINATE, which tells the scheduler the stage must be rolled back and recomputed. Any shuffle output is UNORDERED, so in a chain of two round-robin repartitions the second shuffle's producer reads an UNORDERED input and becomes INDETERMINATE. The DAGScheduler rejects an indeterminate pipelined producer outright, so such a chain failed with PIPELINED_SHUFFLE_UNSUPPORTED (before the previous commit it hung silently instead). A pipelined stage is never retried -- a pipelined task set gets a single attempt -- and is never rolled back and recomputed, so the marking protects nothing there. Exempt pipelined shuffles from it, which lets a chain of round-robin repartitions run. Non-pipelined shuffles are unaffected. Co-authored-by: Isaac --- .../exchange/ShuffleExchangeExec.scala | 12 +++++-- .../streaming/StreamRealTimeModeSuite.scala | 33 +++++++++++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala index 92893a332ba08..2f07702337bb8 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala @@ -512,8 +512,16 @@ object ShuffleExchangeExec { // round-robin function is order sensitive if we don't sort the input. // Stateful partition assignment is order-sensitive when it depends on row visitation order. - val isOrderSensitive = - (isRoundRobin || isNullAwareHashPartitioning) && !SQLConf.get.sortBeforeRepartition + // + // A pipelined shuffle is exempt. Marking the map RDD order-sensitive only serves to make it + // INDETERMINATE when its own input is UNORDERED (see + // MapPartitionsRDD.getOutputDeterministicLevel), which tells the scheduler a retry cannot be + // trusted and the stage must be rolled back and recomputed. A pipelined stage is never + // retried (a pipelined task set gets a single attempt) and never recomputed, and the + // DAGScheduler rejects an indeterminate pipelined producer outright -- so keeping the flag + // would reject a chain of round-robin repartitions rather than protect anything. + val isOrderSensitive = (isRoundRobin || isNullAwareHashPartitioning) && + !SQLConf.get.sortBeforeRepartition && !pipelined if (needToCopyObjectsBeforeShuffle(part)) { newRdd.mapPartitionsWithIndexInternal((_, iter) => { val getPartitionKey = getPartitionKeyExtractor() diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala index 0bda77e8415dc..4f5893fe04dbc 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala @@ -564,6 +564,39 @@ class StreamRealTimeModeWithManualClockSuite extends StreamRealTimeModeManualClo } } + test("pipelined shuffle: a chain of two round-robin repartitions runs in Real-Time Mode") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val inputData = LowLatencyMemoryStream[(String, Int)](2) + // Two round-robin repartitions with a map between them; consecutive repartitions collapse to + // the last one, so the map is what keeps both shuffles in the plan. Both are pipelined, and + // the second one's producer reads the first one's output -- an UNORDERED input. A round-robin + // shuffle is order-sensitive once the deterministic local sort is off, which would escalate + // that producer to INDETERMINATE and get the group rejected; a pipelined shuffle is exempt + // from the order-sensitive marking precisely so this shape runs. + val result = inputData.toDS() + .repartition(3) + .map(row => row) + .repartition(2) + .map(row => row) + .toDF().select($"_1".as("key")) + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, ("a", 1), ("b", 2), ("c", 3)), + StartStream(), + CheckAnswerWithTimeout(60000, "a", "b", "c"), + Execute { q => + val exchanges = q.lastExecution.executedPlan.collect { + case s: ShuffleExchangeExec => s + } + assert(exchanges.size >= 2, s"expected >= 2 shuffle exchanges, got ${exchanges.size}") + assert(exchanges.forall(_.pipelined), + "every exchange in the chain must be pipelined, got: " + + exchanges.map(_.pipelined).mkString(", ")) + }, + StopStream + ) + } + } + test("pipelined shuffle: dedup recovers from a task failure via checkpoint restart") { withTempDir { checkpointDir => // A UDF that throws on demand, placed after the dedup so the failure lands in the pipelined From a810c0160ef752448a0a9167239f232d0da64d8b Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sun, 2 Aug 2026 01:09:40 +0000 Subject: [PATCH 20/24] [SPARK-XXXXX][SS] Always disable sort before repartition in Real-Time Mode The previous guard only set the config when the user had not set it explicitly. Since it defaults to true, a user (or a spark-defaults deployment) that sets it explicitly to true kept the fully-blocking local sort before a round-robin repartition -- which never drains an unbounded stream, so the query hangs forever with no error. That is the exact hazard this change set removes. Override it unconditionally for Real-Time Mode. A query that never produces a row is worse than ignoring a config that only guards against a task retry that cannot happen here. Log a warning when overriding an explicit setting so the operator can see their config is not the one in force. Co-authored-by: Isaac --- .../runtime/MicroBatchExecution.scala | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) 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 f608abd4423a6..39fb0038ad4ae 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 @@ -382,12 +382,20 @@ class MicroBatchExecution( // // Determinism is not needed here: Real-Time Mode does not retry tasks (TaskSetManager caps a // pipelined task set at one attempt, and its group is aborted as a unit rather than - // recomputed), so the hazard the sort guards against does not arise. Only set the config when - // the user has not made an explicit choice. - if (!sparkSessionForStream.conf.contains(SQLConf.SORT_BEFORE_REPARTITION.key)) { - logInfo("Disabling sort before repartition for Real-Time Mode") - sparkSessionForStream.conf.set(SQLConf.SORT_BEFORE_REPARTITION.key, "false") + // recomputed), so the hazard the sort guards against does not arise. + // + // This overrides an explicit `true` as well. Honouring it would reintroduce the hang above, + // and a query that never produces a row is worse than silently ignoring a config that only + // guards against a retry that cannot happen here. Log at warning level when overriding an + // explicit setting so the operator can see their config is not the one in force. + if (sparkSessionForStream.conf.contains(SQLConf.SORT_BEFORE_REPARTITION.key) && + sparkSessionForStream.conf.get(SQLConf.SORT_BEFORE_REPARTITION)) { + logWarning(log"Ignoring ${MDC(LogKeys.CONFIG, SQLConf.SORT_BEFORE_REPARTITION.key)}=true " + + log"for this Real-Time Mode query: the local sort it inserts before a round-robin " + + log"repartition never completes on an unbounded stream. It is not needed because a " + + log"Real-Time Mode task is not retried.") } + sparkSessionForStream.conf.set(SQLConf.SORT_BEFORE_REPARTITION.key, "false") } // Initializing TriggerExecutor relies on `sources`, hence calling this after initializing From c0b715c124af3388322b18af5fcf063f88c8ca28 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sun, 2 Aug 2026 01:09:49 +0000 Subject: [PATCH 21/24] [SPARK-XXXXX][SS] Only mark shuffles on the streaming path as pipelined The marking rule flipped every ShuffleExchangeExec in any plan that contained a RealTimeStreamScanExec anywhere, without checking the exchange was on the streaming path. A plan can hold a static subtree alongside the streaming one: the static side of a broadcast stream-static join is planned in the same physical plan and may carry its own shuffle. That shuffle must materialize normally -- a static side runs to completion rather than streaming -- so marking it pipelined pulled it into the pipelined group and demanded concurrent slots for stages that must instead finish, failing admission with CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT. Mark only shuffles whose subtree reaches the real-time scan. This mirrors the streaming-path detection the operator allowlist already uses, so the rule no longer marks a wider set of exchanges than the allowlist validates. Co-authored-by: Isaac --- .../runtime/IncrementalExecution.scala | 33 +++++++++- .../streaming/StreamRealTimeModeSuite.scala | 64 +++++++++++++++++++ 2 files changed, 94 insertions(+), 3 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala index 18644e0e780c2..be5d259261a37 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala @@ -699,11 +699,38 @@ class IncrementalExecution( if (!isRealTimeMode) { plan } else { - plan.transformUp { - case s: ShuffleExchangeExec if !s.pipelined => s.copy(pipelined = true) - } + markStreamingPath(plan)._1 } } + + /** + * Marks the shuffles that are on the streaming path -- those whose subtree reaches a + * [[RealTimeStreamScanExec]] -- and returns the rewritten plan along with whether this + * subtree reaches one. + * + * A plan can hold a static subtree alongside the streaming one: the static side of a + * broadcast stream-static join is planned in the same physical plan and may contain its own + * shuffle. That shuffle materializes normally and is not part of the pipelined group -- and + * cannot be, since a static side runs to completion rather than streaming. Marking it + * pipelined would pull it into the group and demand slots for stages that must instead + * finish, which fails admission (CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT). This mirrors the + * streaming-path detection the operator allowlist uses (RealTimeModeAllowlist), which only + * inspects nodes whose subtree reaches the real-time scan; marking a wider set than the + * allowlist checks would flip shuffles it never validated. + */ + private def markStreamingPath(plan: SparkPlan): (SparkPlan, Boolean) = plan match { + case rts: RealTimeStreamScanExec => (rts, true) + case p if p.children.isEmpty => (p, false) + case p => + val results = p.children.map(markStreamingPath) + val onStreamingPath = results.exists(_._2) + val newPlan = p.withNewChildren(results.map(_._1)) + newPlan match { + case s: ShuffleExchangeExec if onStreamingPath && !s.pipelined => + (s.copy(pipelined = true), true) + case other => (other, onStreamingPath) + } + } } override def preparations: Seq[Rule[SparkPlan]] = diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala index 4f5893fe04dbc..d0fd39db9e69c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala @@ -27,6 +27,7 @@ import org.scalatest.concurrent.PatienceConfiguration.Timeout import org.apache.spark.{SparkException, SparkIllegalArgumentException, SparkIllegalStateException, TaskContext} import org.apache.spark.scheduler.{SparkListener, SparkListenerStageCompleted, SparkListenerStageSubmitted} +import org.apache.spark.sql.execution.datasources.v2.RealTimeStreamScanExec import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.execution.streaming.RealTimeTrigger import org.apache.spark.sql.execution.streaming.runtime.{MemoryStream, StreamExecution} @@ -214,6 +215,48 @@ class StreamRealTimeModeSuite extends StreamRealTimeModeSuiteBase { ) } + test("pipelined shuffle: a static-side shuffle is not marked pipelined") { + // A broadcast stream-static join can carry a shuffle on its STATIC side, in a subtree with no + // RealTimeStreamScanExec. That shuffle must materialize normally: a static side runs to + // completion rather than streaming, so pulling it into the pipelined group would demand + // concurrent slots for a stage that must instead finish, failing admission. Only shuffles on + // the streaming path are marked, matching what the operator allowlist actually validates. + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10485760") { + val staticData = spark.range(0, 3).selectExpr("id as sk", "id * 10 as sv") + .repartition(3, $"sk") + val inputData = LowLatencyMemoryStream[(String, Int)](2) + val streamDf = inputData.toDF().select($"_1".as("key"), $"_2".cast("long").as("sk")) + val result = streamDf + .repartition(2, $"sk") + .join(broadcast(staticData), Seq("sk"), "left") + .select($"key") + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, ("a", 1), ("b", 2)), + StartStream(), + CheckAnswerWithTimeout(60000, "a", "b"), + Execute { q => + val exchanges = q.lastExecution.executedPlan.collect { + case s: ShuffleExchangeExec => s + } + // The streaming-side repartition is pipelined; a static-side shuffle, if it survives + // into this plan, must not be. + val streamingSide = exchanges.filter(_.exists { + case _: RealTimeStreamScanExec => true + case _ => false + }) + assert(streamingSide.nonEmpty, "expected a shuffle on the streaming path") + assert(streamingSide.forall(_.pipelined), + "a streaming-path shuffle must be pipelined") + assert(exchanges.filterNot(streamingSide.contains).forall(!_.pipelined), + "a static-side shuffle must not be marked pipelined") + }, + StopStream + ) + } + } + test("multiple broadcast joins with the same static table run in Real-Time Mode") { // Two broadcast joins against the SAME static table -- exchange reuse collapses the second // broadcast into a ReusedExchangeExec. This is the supported reuse shape in Real-Time Mode: a @@ -597,6 +640,27 @@ class StreamRealTimeModeWithManualClockSuite extends StreamRealTimeModeManualClo } } + test("pipelined shuffle: an explicit sortBeforeRepartition=true does not stall the query") { + // The deterministic local sort before a round-robin repartition never drains an unbounded + // stream, so Real-Time Mode overrides this config even when it is set explicitly -- honouring + // it would hang the query forever. See MicroBatchExecution. + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.SORT_BEFORE_REPARTITION.key -> "true") { + val inputData = LowLatencyMemoryStream[(String, Int)](2) + val result = inputData.toDF().select($"_1".as("key")) + .repartition(4) + .dropDuplicates("key") + .select($"key") + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, ("a", 1), ("b", 1), ("c", 1), ("a", 2)), + StartStream(), + CheckAnswerWithTimeout(60000, "a", "b", "c"), + StopStream + ) + } + } + test("pipelined shuffle: dedup recovers from a task failure via checkpoint restart") { withTempDir { checkpointDir => // A UDF that throws on demand, placed after the dedup so the failure lands in the pipelined From 685ddeb7db8a30a63b52bc57464cb184152c456f Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sun, 2 Aug 2026 01:16:07 +0000 Subject: [PATCH 22/24] [SPARK-XXXXX][SS][TESTS] Make the Real-Time Mode shuffle tests assert what they claim Two tests could pass without exercising the behaviour they describe: - The co-scheduling tests counted every stage in the shared SparkContext, so an unrelated concurrent stage was enough to satisfy `maxConcurrentStages >= 2` even if this query's producer and consumer had run one after the other. Count only stages belonging to the query under test, identified by the query-id property StreamExecution tags its jobs with. - The reused-broadcast test asserted only the join answer, so it would pass whether or not exchange reuse occurred. Assert the plan actually contains a ReusedExchangeExec and no shuffle exchange, which is the shape the test is about. Co-authored-by: Isaac --- .../streaming/StreamRealTimeModeSuite.scala | 52 ++++++++++++++++--- 1 file changed, 45 insertions(+), 7 deletions(-) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala index d0fd39db9e69c..36fa816f55934 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala @@ -26,9 +26,9 @@ import scala.concurrent.duration.Duration import org.scalatest.concurrent.PatienceConfiguration.Timeout import org.apache.spark.{SparkException, SparkIllegalArgumentException, SparkIllegalStateException, TaskContext} -import org.apache.spark.scheduler.{SparkListener, SparkListenerStageCompleted, SparkListenerStageSubmitted} +import org.apache.spark.scheduler.{SparkListener, SparkListenerJobStart, SparkListenerStageCompleted, SparkListenerStageSubmitted} import org.apache.spark.sql.execution.datasources.v2.RealTimeStreamScanExec -import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec} import org.apache.spark.sql.execution.streaming.RealTimeTrigger import org.apache.spark.sql.execution.streaming.runtime.{MemoryStream, StreamExecution} import org.apache.spark.sql.execution.streaming.sources.{ContinuousMemorySink, LowLatencyMemoryStream} @@ -269,7 +269,17 @@ class StreamRealTimeModeSuite extends StreamRealTimeModeSuiteBase { testStream(df2, OutputMode.Update, Map.empty, new ContinuousMemorySink())( AddData(inputData, 1, 2, 3), StartStream(), - CheckAnswerWithTimeout(30000, (1, "a", "a"), (2, "b", "b"), (3, "c", "c")) + CheckAnswerWithTimeout(30000, (1, "a", "a"), (2, "b", "b"), (3, "c", "c")), + Execute { q => + val plan = q.lastExecution.executedPlan + // Verify the shape this test is about, not just the answer: the second broadcast is reused, + // and the marking rule left it alone because it only marks shuffle exchanges. + assert(plan.exists(_.isInstanceOf[ReusedExchangeExec]), + s"expected the second broadcast to be reused, got:\n$plan") + assert(plan.collect { case s: ShuffleExchangeExec => s }.isEmpty, + "a broadcast stream-static join should introduce no shuffle exchange") + }, + StopStream ) } } @@ -484,10 +494,24 @@ class StreamRealTimeModeWithManualClockSuite extends StreamRealTimeModeManualClo // schedule never exceeds one running stage at a time; >= 2 proves genuine co-scheduling. val runningStages = ConcurrentHashMap.newKeySet[Int]() val maxConcurrentStages = new AtomicInteger(0) + val queryStageIds = ConcurrentHashMap.newKeySet[Int]() + // Count only stages belonging to the query under test. The suite shares one SparkContext, so + // an unrelated concurrent stage would otherwise satisfy the co-scheduling assertion below even + // if this query's producer and consumer actually ran one after the other. StreamExecution sets + // the job group to the query's runId, which the job-start event carries. val listener = new SparkListener { + override def onJobStart(e: SparkListenerJobStart): Unit = { + // StreamExecution tags every streaming job with the query id; a batch/unrelated job has + // no such property, so this keeps unrelated stages out of the count. + if (e.properties.getProperty(StreamExecution.QUERY_ID_KEY) != null) { + e.stageIds.foreach(queryStageIds.add(_)) + } + } override def onStageSubmitted(e: SparkListenerStageSubmitted): Unit = { - runningStages.add(e.stageInfo.stageId) - maxConcurrentStages.accumulateAndGet(runningStages.size(), Math.max) + if (queryStageIds.contains(e.stageInfo.stageId)) { + runningStages.add(e.stageInfo.stageId) + maxConcurrentStages.accumulateAndGet(runningStages.size(), Math.max) + } } override def onStageCompleted(e: SparkListenerStageCompleted): Unit = { runningStages.remove(e.stageInfo.stageId) @@ -564,10 +588,24 @@ class StreamRealTimeModeWithManualClockSuite extends StreamRealTimeModeManualClo // group scenario: more than one pipelined shuffle in a single Real-Time Mode job. val runningStages = ConcurrentHashMap.newKeySet[Int]() val maxConcurrentStages = new AtomicInteger(0) + val queryStageIds = ConcurrentHashMap.newKeySet[Int]() + // Count only stages belonging to the query under test. The suite shares one SparkContext, so + // an unrelated concurrent stage would otherwise satisfy the co-scheduling assertion below even + // if this query's producer and consumer actually ran one after the other. StreamExecution sets + // the job group to the query's runId, which the job-start event carries. val listener = new SparkListener { + override def onJobStart(e: SparkListenerJobStart): Unit = { + // StreamExecution tags every streaming job with the query id; a batch/unrelated job has + // no such property, so this keeps unrelated stages out of the count. + if (e.properties.getProperty(StreamExecution.QUERY_ID_KEY) != null) { + e.stageIds.foreach(queryStageIds.add(_)) + } + } override def onStageSubmitted(e: SparkListenerStageSubmitted): Unit = { - runningStages.add(e.stageInfo.stageId) - maxConcurrentStages.accumulateAndGet(runningStages.size(), Math.max) + if (queryStageIds.contains(e.stageInfo.stageId)) { + runningStages.add(e.stageInfo.stageId) + maxConcurrentStages.accumulateAndGet(runningStages.size(), Math.max) + } } override def onStageCompleted(e: SparkListenerStageCompleted): Unit = { runningStages.remove(e.stageInfo.stageId) From f14e9993315c1ae6f8f78d5c4bfa56c95f9b13d7 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sun, 2 Aug 2026 01:52:40 +0000 Subject: [PATCH 23/24] [SPARK-XXXXX][SS] Do not mark a range-partitioned shuffle as pipelined Building a RangePartitioner runs a separate job that samples the input to compute range bounds, which cannot complete on an unbounded stream. Marking such a shuffle pipelined would pull that sampling job into a pipelined group, so leave it as a normal materializing shuffle instead. This does not make repartitionByRange work in Real-Time Mode -- the sampling job cannot terminate on an unbounded source either way -- it only stops the rule from labelling a shuffle the pipelined path cannot serve. Co-authored-by: Isaac --- .../runtime/IncrementalExecution.scala | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala index be5d259261a37..d6b97c629b485 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala @@ -30,6 +30,7 @@ import org.apache.spark.sql.catalyst.QueryPlanningTracker import org.apache.spark.sql.catalyst.analysis.WidenStatefulOperatorAttributeNullability import org.apache.spark.sql.catalyst.expressions.{CurrentBatchTimestamp, ExpressionWithRandomSeed} import org.apache.spark.sql.catalyst.plans.logical._ +import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, RangePartitioning} import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.trees.TreePattern._ import org.apache.spark.sql.classic.SparkSession @@ -726,11 +727,26 @@ class IncrementalExecution( val onStreamingPath = results.exists(_._2) val newPlan = p.withNewChildren(results.map(_._1)) newPlan match { - case s: ShuffleExchangeExec if onStreamingPath && !s.pipelined => + case s: ShuffleExchangeExec + if onStreamingPath && !s.pipelined && canBePipelined(s.outputPartitioning) => (s.copy(pipelined = true), true) case other => (other, onStreamingPath) } } + + /** + * Whether a shuffle with this partitioning can be pipelined. + * + * `RangePartitioning` cannot: building its `RangePartitioner` runs a separate job that samples + * the input to compute range bounds (see `ShuffleExchangeExec.prepareShuffleDependency`), which + * cannot complete on an unbounded stream. Leaving it unmarked keeps it a normal materializing + * shuffle rather than silently pulling a sampling job into a pipelined group. Such a query does + * not run in Real-Time Mode either way -- this only avoids mislabelling the shuffle. + */ + private def canBePipelined(partitioning: Partitioning): Boolean = partitioning match { + case _: RangePartitioning => false + case _ => true + } } override def preparations: Seq[Rule[SparkPlan]] = From e1ea484b7e5c366e24f0e8c48f673bdad05da872 Mon Sep 17 00:00:00 2001 From: Jerry Peng Date: Sun, 2 Aug 2026 03:02:07 +0000 Subject: [PATCH 24/24] [SPARK-XXXXX][SS] Reject a range-partitioned shuffle in Real-Time Mode Allowlisting ShuffleExchangeExec by class name admits every partitioning, but only some can run in Real-Time Mode. Building a RangePartitioner runs a separate job that samples the input to compute range bounds, and that job cannot complete while the source keeps producing. Before this change set the allowlist rejected every shuffle, so a repartitionByRange failed fast; admitting shuffles turned that into a query that stalls indefinitely with no error. Reject a range-partitioned shuffle at admission so it fails fast with the usual allowlist error again. Co-authored-by: Isaac --- .../runtime/RealTimeModeAllowlist.scala | 21 +++++++++++++++- .../StreamRealTimeModeAllowlistSuite.scala | 25 +++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala index 047ea3de4298d..d89117302f97c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala @@ -20,8 +20,10 @@ package org.apache.spark.sql.execution.streaming.runtime import org.apache.spark.SparkIllegalArgumentException import org.apache.spark.internal.{Logging, LogKeys, MessageWithContext} import org.apache.spark.sql.connector.catalog.Table +import org.apache.spark.sql.catalyst.plans.physical.RangePartitioning import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.datasources.v2.RealTimeStreamScanExec +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.execution.streaming.operators.stateful._ object RealTimeModeAllowlist extends Logging { @@ -130,12 +132,29 @@ object RealTimeModeAllowlist extends Logging { collectNodesWhoseSubtreeHasRTS(root)._2 } + /** + * Whether this operator is supported in Real-Time Mode. + * + * Membership in `allowedOperators` is the general rule, but a shuffle needs a second look: only + * some partitionings can run in Real-Time Mode. A `RangePartitioning` cannot, because building + * its `RangePartitioner` runs a separate job that samples the input to compute range bounds (see + * `ShuffleExchangeExec.prepareShuffleDependency`), and that job cannot complete while the source + * keeps producing. Rejecting it here fails the query fast with the usual allowlist error instead + * of letting it stall on a sampling job that never finishes. + */ + private def isAllowed(node: SparkPlan): Boolean = { + allowedOperators.contains(node.getClass.getName) && (node match { + case e: ShuffleExchangeExec => !e.outputPartitioning.isInstanceOf[RangePartitioning] + case _ => true + }) + } + def checkAllowedPhysicalOperator(operator: SparkPlan, throwException: Boolean): Unit = { val nodesToCheck = collectRealtimeNodes(operator) val violations = nodesToCheck .collect { case node => - if (allowedOperators.contains(node.getClass.getName)) { + if (isAllowed(node)) { None } else { Some(node.getClass.getName) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala index bad7cd2cf695f..de4af97656a69 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala @@ -106,6 +106,31 @@ class StreamRealTimeModeAllowlistSuite extends StreamRealTimeModeE2ESuiteBase { } } + // A repartitionByRange produces a range-partitioned shuffle. Building its RangePartitioner runs a + // separate job that samples the input to compute range bounds, which cannot complete while the + // source keeps producing, so it must be rejected up front rather than stalling the query. + test("range-partitioned shuffle not allowed") { + val inputData = LowLatencyMemoryStream[Int](2) + + val df = inputData.toDF() + .select(col("value").as("key")) + .repartitionByRange(3, col("key")) + .select(col("key")) + + val query = runStreamingQuery("range_shuffle_allowlist", df) + + eventually(timeout(60.seconds)) { + checkError( + exception = query.exception.get.getCause.asInstanceOf[SparkIllegalArgumentException], + condition = "STREAMING_REAL_TIME_MODE.OPERATOR_OR_SINK_NOT_IN_ALLOWLIST", + parameters = Map( + "errorType" -> "operator", + "message" -> "org.apache.spark.sql.execution.exchange.ShuffleExchangeExec is" + ) + ) + } + } + // A stateful aggregation is still rejected under RTM because HashAggregateExec is not in the // allowlist, even though the pipelined shuffle and the StateStore operators now are. When RTM // supports aggregation (SPARK-54236), remove this test.