cloud-fan commented on code in PR #57692:
URL: https://github.com/apache/spark/pull/57692#discussion_r3709557379
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala:
##########
@@ -663,7 +664,91 @@ 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.
+ *
+ * Marks every eligible shuffle exchange on the streaming path, 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). There is no shuffle-count restriction. An exchange
whose subtree does not
+ * reach the real-time scan is skipped -- the static side of a broadcast
stream-static join must
+ * materialize, because it runs to completion rather than streaming. A
partitioning the pipelined
+ * path cannot serve, such as range partitioning, is rejected up front by
RealTimeModeAllowlist
+ * rather than being handled here.
+ *
+ * The walk 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 when the query starts (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 rejected
by the DAGScheduler
+ * (checkPipelinedGroupsSupportedInRDDGraph). Note that check runs in
handleJobSubmitted, so it
+ * rejects the batch's job rather than the query: such a query fails the
same way on every batch
+ * instead of failing once when it is planned. Marking is not what makes the
shape unsupported, so
Review Comment:
This reverses the cause of the rejection: the scheduler rejects fan-out
specifically for a `PipelinedShuffleDependency`, while a regular materialized
shuffle supports it. Please describe fan-out as a limitation activated by the
pipelined path; a plan-time guard would improve only when the error is reported.
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala:
##########
@@ -393,4 +466,330 @@ 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 queryStageIds = ConcurrentHashMap.newKeySet[Int]()
+ // Count only stages belonging to the query under test. The suite shares
one SparkContext, so a
+ // stage from any other streaming query would otherwise satisfy the
co-scheduling assertion
+ // below even if this query's producer and consumer actually ran one after
the other. The id is
+ // captured from the query once it is running, and every job is matched
against it.
+ val queryId = new AtomicReference[String](null)
+ val listener = new SparkListener {
+ override def onJobStart(e: SparkListenerJobStart): Unit = {
+ // StreamExecution tags every streaming job with its query id.
+ val id = queryId.get()
+ if (id != null &&
e.properties.getProperty(StreamExecution.QUERY_ID_KEY) == id) {
+ e.stageIds.foreach(queryStageIds.add(_))
+ }
+ }
+ override def onStageSubmitted(e: SparkListenerStageSubmitted): Unit = {
+ 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)
+ }
+ }
+ 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(),
+ // Record the id before any batch is awaited, so the listener
attributes this query's jobs
+ // from the first one.
+ Execute(q => queryId.set(q.id.toString)),
Review Comment:
`StartStream` can launch the first job before this action sets `queryId`, so
that job never enters `queryStageIds`. This is especially flaky in the
two-shuffle test, which asserts before running another batch. Set the ID, then
trigger a dedicated observed batch in both listener-based tests.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
