zhongyu09 commented on a change in pull request #31167:
URL: https://github.com/apache/spark/pull/31167#discussion_r560637456
##########
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:
> 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.
submitMapState in DAGScheduler is introduced in SPARK-9851 just for AQE
propose. I also think it can be used in normal query. I guess the reason we
don't is there is less benefit as the complex it introduced. The architecture
of sql execution should be changed just like AQE, one job will be break to many
job, etc..
> You could say in non-AQE currently shuffle map tasks are submitted after
the broadcast finishes. But not we always wait for the broadcast.
It seems like the same thing. But we can make it better in AQE. If we don't
like the solution, we can use the initial one as partial fix and find a good
solution, as discussed with @cloud-fan .
----------------------------------------------------------------
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:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
