set correct execution Id for broadcast query stage (apache#50)

mccheah · Jan 14, 2019 · f3b7dee · f3b7dee
1 parent fff789a
commit f3b7dee
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/QueryStage.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/QueryStage.scala
@@ -53,16 +53,21 @@ abstract class QueryStage extends UnaryExecNode {
    * blocking on one child stage.
    */
   def executeChildStages(): Unit = {
+    val executionId = sqlContext.sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)
+
     // Handle broadcast stages
     val broadcastQueryStages: Seq[BroadcastQueryStage] = child.collect {
       case bqs: BroadcastQueryStageInput => bqs.childStage
     }
     val broadcastFutures = broadcastQueryStages.map { queryStage =>
-      Future { queryStage.prepareBroadcast() }(QueryStage.executionContext)
+      Future {
+        SQLExecution.withExecutionId(sqlContext.sparkContext, executionId) {
+          queryStage.prepareBroadcast()
+        }
+      }(QueryStage.executionContext)
     }
 
     // Submit shuffle stages
-    val executionId = sqlContext.sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)
     val shuffleQueryStages: Seq[ShuffleQueryStage] = child.collect {
       case sqs: ShuffleQueryStageInput => sqs.childStage
     }