jerrypeng commented on code in PR #57692:
URL: https://github.com/apache/spark/pull/57692#discussion_r3717345462
##########
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:
I took the second half of your suggestion (a dedicated observed batch) but
not the first, because the id can't be set before StartStream.
--
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]
