viirya commented on a change in pull request #31167:
URL: https://github.com/apache/spark/pull/31167#discussion_r560413130
##########
File path:
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala
##########
@@ -190,7 +191,36 @@ case class AdaptiveSparkPlanExec(
executionId.foreach(onUpdatePlan(_, result.newStages.map(_.plan)))
// Start materialization of all new stages and fail fast if any
stages failed eagerly
- result.newStages.foreach { stage =>
+
+ // SPARK-33933: we should materialize broadcast stages first and
wait the
+ // materialization finish before materialize other stages, to avoid
waiting
+ // for broadcast tasks to be scheduled and leading to broadcast
timeout.
+ val broadcastMaterializationFutures = result.newStages
+ .filter(_.isInstanceOf[BroadcastQueryStageExec])
+ .map { stage =>
+ var future: Future[Any] = null
+ try {
+ future = stage.materialize()
+ future.onComplete { res =>
+ if (res.isSuccess) {
+ events.offer(StageSuccess(stage, res.get))
+ } else {
+ events.offer(StageFailure(stage, res.failed.get))
+ }
+ }(AdaptiveSparkPlanExec.executionContext)
+ } catch {
+ case e: Throwable =>
+ cleanUpAndThrowException(Seq(e), Some(stage.id))
+ }
+ future
+ }
+
+ // Wait for the materialization of all broadcast stages finish
Review comment:
Ok, I see your point.
> I can ensure there's enough resources to run broadcast and shuffle in
parallel. Actually, there haven't come to TaskScheduler, "sc.runJob" is called
after the broadcast finished.
In short, `submitMapState` can be used to submit a job too. Although seems
it is only called now in AQE. I'm surprised that we don't do it similarly in
normal query.
> In non-AQE, we always wait the broadcast finish before submit shuffle map
tasks.
I think this isn't true because we don't forcibly wait broadcast to be
finished like the change here. Technically, Spark doesn't restrict a shuffle to
be run in parallel with a broadcast. An operator can call `submitMapState` to
trigger a shuffle map stage if it wants, and the map stage can be run in
parallel with broadcast.
You could say in non-AQE currently shuffle map tasks are submitted after the
broadcast finishes. But not we always wait for the broadcast.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
