Add public finalPhysicalPlan that uses withFinalPlanUpdate

apache · Nov 8, 2022 · 52b7bb6 · 52b7bb6
1 parent 769e183
commit 52b7bb6
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 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
@@ -222,7 +222,9 @@ case class AdaptiveSparkPlanExec(
       .map(_.toLong).filter(SQLExecution.getQueryExecution(_) eq context.qe)
   }
 
-  private[sql] def getFinalPhysicalPlan(): SparkPlan = lock.synchronized {
+  def finalPhysicalPlan: SparkPlan = withFinalPlanUpdate(identity)
+
+  private def getFinalPhysicalPlan(): SparkPlan = lock.synchronized {
     if (isFinalPlan) return currentPhysicalPlan
 
     // In case of this adaptive plan being executed out of `withActive` scoped functions, e.g.,

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileFormatWriter.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileFormatWriter.scala
@@ -192,7 +192,7 @@ object FileFormatWriter extends Logging {
     // SPARK-40588: plan may contain an AdaptiveSparkPlanExec, which does not know
     // its final plan's ordering, so we have to materialize that plan first
     def materializeAdaptiveSparkPlan(plan: SparkPlan): SparkPlan = plan match {
-      case a: AdaptiveSparkPlanExec => a.getFinalPhysicalPlan()
+      case a: AdaptiveSparkPlanExec => a.finalPhysicalPlan
       case p: SparkPlan => p.withNewChildren(p.children.map(materializeAdaptiveSparkPlan))
     }
     val materializedPlan = materializeAdaptiveSparkPlan(empty2NullPlan)