This is an automated email from the ASF dual-hosted git repository.
dongjoon pushed a commit to branch branch-3.0
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-3.0 by this push:
new e1ed80a [SPARK-31473][SQL] AQE should set active session during
execution
e1ed80a is described below
commit e1ed80a134a2d613a9221a15847cad3274cd6f2b
Author: Maryann Xue <maryann....@gmail.com>
AuthorDate: Sat Apr 18 00:08:36 2020 -0700
[SPARK-31473][SQL] AQE should set active session during execution
### What changes were proposed in this pull request?
AQE creates new SparkPlan nodes during execution. This PR makes sure that
the active session is set correctly during this process and AQE execution is
not disrupted by external session change.
### Why are the changes needed?
To prevent potential errors. If not changed, the physical plans generated
by AQE would have the wrong SparkSession or even null SparkSession, which could
lead to NPE.
### Does this PR introduce any user-facing change?
No.
### How was this patch tested?
Added UT.
Closes #28247 from maryannxue/aqe-activesession.
Authored-by: Maryann Xue <maryann....@gmail.com>
Signed-off-by: Dongjoon Hyun <dongj...@apache.org>
(cherry picked from commit 6198f384054e7f86521891ceeb1a231f449a16a8)
Signed-off-by: Dongjoon Hyun <dongj...@apache.org>
---
.../sql/execution/adaptive/AdaptiveSparkPlanExec.scala | 9 +++++++--
.../sql/execution/adaptive/AdaptiveQueryExecSuite.scala | 13 ++++++++++++-
2 files changed, 19 insertions(+), 3 deletions(-)
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala
index b54a32f..2b46724 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala
@@ -139,7 +139,12 @@ case class AdaptiveSparkPlanExec(
}
private def getFinalPhysicalPlan(): SparkPlan = lock.synchronized {
- if (!isFinalPlan) {
+ if (isFinalPlan) return currentPhysicalPlan
+
+ // In case of this adaptive plan being executed out of `withActive` scoped
functions, e.g.,
+ // `plan.queryExecution.rdd`, we need to set active session here as new
plan nodes can be
+ // created in the middle of the execution.
+ context.session.withActive {
// Subqueries do not have their own execution IDs and therefore rely on
the main query to
// update UI.
val executionId = Option(context.session.sparkContext.getLocalProperty(
@@ -225,8 +230,8 @@ case class AdaptiveSparkPlanExec(
isFinalPlan = true
executionId.foreach(onUpdatePlan(_, Seq(currentPhysicalPlan)))
logOnLevel(s"Final plan: $currentPhysicalPlan")
+ currentPhysicalPlan
}
- currentPhysicalPlan
}
override def executeCollect(): Array[InternalRow] = {
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
index 64dd9aa..6da510f 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
@@ -23,13 +23,14 @@ import java.net.URI
import org.apache.log4j.Level
import org.apache.spark.scheduler.{SparkListener, SparkListenerEvent,
SparkListenerJobStart}
-import org.apache.spark.sql.{QueryTest, Row}
+import org.apache.spark.sql.{QueryTest, Row, SparkSession}
import org.apache.spark.sql.execution.{ReusedSubqueryExec, ShuffledRowRDD,
SparkPlan}
import
org.apache.spark.sql.execution.adaptive.OptimizeLocalShuffleReader.LOCAL_SHUFFLE_READER_DESCRIPTION
import org.apache.spark.sql.execution.command.DataWritingCommandExec
import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec,
Exchange, ReusedExchangeExec}
import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec,
BuildRight, SortMergeJoinExec}
import
org.apache.spark.sql.execution.ui.SparkListenerSQLAdaptiveExecutionUpdate
+import org.apache.spark.sql.functions._
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
import org.apache.spark.sql.types.{IntegerType, StructType}
@@ -805,4 +806,14 @@ class AdaptiveQueryExecSuite
assert(plan.asInstanceOf[DataWritingCommandExec].child.isInstanceOf[AdaptiveSparkPlanExec])
}
}
+
+ test("AQE should set active session during execution") {
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") {
+ val df = spark.range(10).select(sum('id))
+
assert(df.queryExecution.executedPlan.isInstanceOf[AdaptiveSparkPlanExec])
+ SparkSession.setActiveSession(null)
+ checkAnswer(df, Seq(Row(45)))
+ SparkSession.setActiveSession(spark) // recover the active session.
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@spark.apache.org
For additional commands, e-mail: commits-h...@spark.apache.org
