cloud-fan commented on code in PR #57692:
URL: https://github.com/apache/spark/pull/57692#discussion_r3702359158
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala:
##########
@@ -663,7 +665,96 @@ 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 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 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
+ * a plan-time guard here would only improve the failure mode, not the
outcome.
+ */
+ object MarkPipelinedShuffleForRealTimeMode extends Rule[SparkPlan] {
+ override def apply(plan: SparkPlan): SparkPlan = {
+ val isRealTimeMode = plan.exists(_.isInstanceOf[RealTimeStreamScanExec])
+ if (!isRealTimeMode) {
+ plan
+ } else {
+ 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 &&
canBePipelined(s.outputPartitioning) =>
+ (s.copy(pipelined = true), true)
Review Comment:
Preserve the original exchange's tree-node tags when setting
`pipelined=true`. This case-class copy bypasses `TreeNode.copyTagsFrom`, so the
marked exchange loses metadata such as its logical link; use a tag-preserving
tree operation or explicitly copy the tags onto the replacement.
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala:
##########
@@ -393,4 +466,322 @@ 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
Review Comment:
Filter these events by this query's actual ID or run ID. Checking only that
`QUERY_ID_KEY` is present admits stages from any concurrent streaming query, so
`maxConcurrentStages` can reach two even when this query's stages run
sequentially. Please apply the same fix to the listener near the two-shuffle
test.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala:
##########
@@ -663,7 +665,96 @@ 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 shuffle exchange, so a plan with several pipelined shuffles
in a chain (e.g. two
Review Comment:
Please narrow this to every eligible shuffle on the streaming path. The rule
intentionally skips static-side and range-partitioned exchanges, so the
universal wording contradicts the implementation.
--
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]
