/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.spark.sql.execution

import java.io.{BufferedWriter, OutputStreamWriter}
import java.util.UUID
import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong}
import javax.annotation.concurrent.GuardedBy

import scala.util.control.NonFatal

import org.apache.hadoop.fs.Path

import org.apache.spark.{SparkContext, SparkException}
import org.apache.spark.internal.LogKeys.EXTENDED_EXPLAIN_GENERATOR
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.{AnalysisException, ExtendedExplainGenerator, Row}
import org.apache.spark.sql.catalyst.{InternalRow, QueryPlanningTracker}
import org.apache.spark.sql.catalyst.analysis.{Analyzer, LazyExpression, NameParameterizedQuery, UnsupportedOperationChecker}
import org.apache.spark.sql.catalyst.expressions.codegen.ByteCodeStats
import org.apache.spark.sql.catalyst.plans.QueryPlan
import org.apache.spark.sql.catalyst.plans.logical.{AppendData, Command, CommandResult, CompoundBody, CreateTableAsSelect, LogicalPlan, OverwriteByExpression, OverwritePartitionsDynamic, ReplaceTableAsSelect, ReturnAnswer, Union, UnresolvedWith, WithCTE}
import org.apache.spark.sql.catalyst.rules.{PlanChangeLogger, Rule}
import org.apache.spark.sql.catalyst.transactions.TransactionUtils
import org.apache.spark.sql.catalyst.util.StringUtils.PlanStringConcat
import org.apache.spark.sql.catalyst.util.truncatedString
import org.apache.spark.sql.classic.SparkSession
import org.apache.spark.sql.connector.catalog.LookupCatalog
import org.apache.spark.sql.connector.catalog.transactions.Transaction
import org.apache.spark.sql.execution.SQLExecution.EXECUTION_ROOT_ID_KEY
import org.apache.spark.sql.execution.adaptive.{AdaptiveExecutionContext, InsertAdaptiveSparkPlan}
import org.apache.spark.sql.execution.bucketing.{CoalesceBucketsInJoin, DisableUnnecessaryBucketedScan}
import org.apache.spark.sql.execution.datasources.v2.{TransactionalExec, V2TableRefreshUtil}
import org.apache.spark.sql.execution.dynamicpruning.PlanDynamicPruningFilters
import org.apache.spark.sql.execution.exchange.EnsureRequirements
import org.apache.spark.sql.execution.reuse.ReuseExchangeAndSubquery
import org.apache.spark.sql.execution.streaming.checkpointing.OffsetSeqMetadata
import org.apache.spark.sql.execution.streaming.runtime.{IncrementalExecution, WatermarkPropagator}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.scripting.SqlScriptingExecution
import org.apache.spark.sql.streaming.OutputMode
import org.apache.spark.util.{LazyTry, Utils, UUIDv7Generator}
import org.apache.spark.util.ArrayImplicits._

/**
 * The primary workflow for executing relational queries using Spark.  Designed to allow easy
 * access to the intermediate phases of query execution for developers.
 *
 * While this is not a public class, we should avoid changing the function names for the sake of
 * changing them, because a lot of developers use the feature for debugging.
 */
class QueryExecution(
    val sparkSession: SparkSession,
    val logical: LogicalPlan,
    val tracker: QueryPlanningTracker = new QueryPlanningTracker,
    val mode: CommandExecutionMode.Value = CommandExecutionMode.ALL,
    val shuffleCleanupModeOpt: Option[ShuffleCleanupMode] = None,
    val refreshPhaseEnabled: Boolean = true,
    val queryId: UUID = UUIDv7Generator.generate(),
    // When a transaction is active, callers creating nested QueryExecution instances MUST pass
    // the enclosing QueryExecution's analyzer here to propagate the transaction context.
    // Omitting it causes the nested QE to use sessionState.analyzer, which has no knowledge
    // of the transaction and will load tables outside the transaction's catalog scope.
    val analyzerOpt: Option[Analyzer] = None) extends LookupCatalog {

  val id: Long = QueryExecution.nextExecutionId

  // Used by SQLExecution to determine whether to use the existing queryId or generate a new one.
  private[sql] val firstExecution = new AtomicBoolean(true)

  // TODO: Move the planner an optimizer into here from SessionState.
  protected def planner = sparkSession.sessionState.planner

  protected val catalogManager = sparkSession.sessionState.catalogManager

  /**
   * Check whether the query represented by this QueryExecution is a SQL script.
   * @return True if the query is a SQL script, False otherwise.
   */
  def isSqlScript: Boolean = QueryExecution.isUnresolvedPlanSqlScript(logical)

  lazy val isLazyAnalysis: Boolean = {
    // Only check the main query as subquery expression can be resolved now with the main query.
    logical.exists(_.expressions.exists(_.exists(_.isInstanceOf[LazyExpression])))
  }


  // 1. At the pre-Analyzed plan we look for nodes that implement the TransactionalWrite trait.
  //    When a plan contains such a node we initiate a transaction. Note, we should never start
  //    a transaction for operations that are not executed, e.g. EXPLAIN.
  // 2. Create an analyzer clone with a transaction aware Catalog Manager. The latter is the
  //    narrow waist of all catalog accesses, and it is also the transaction context carrier.
  //    This is then passed to all rules during analysis that need to check the catalog. Rules
  //    that are specifically interested in transactionality can access the transaction directly
  //    from the Catalog Manager. The transaction catalog, is potentially the place where connectors
  //    should keep state about the reads (tables+predicates) that occurred during the transaction.
  // 3. The analyzer instance is passed to nested Query Execution instances. These need to respect
  //    the open transaction instead of creating their own.
  private val lazyTransactionOpt = LazyTry {
    // Always inherit an active transaction from the outer analyzer, regardless of mode.
    analyzerOpt.flatMap(_.catalogManager.transaction).orElse {
      // Only begin a new transaction for outer QEs that lead to execution.
      if (mode != CommandExecutionMode.SKIP) {
        val catalog = logical match {
          case UnresolvedWith(TransactionalWrite(c), _, _) => Some(c)
          case TransactionalWrite(c) => Some(c)
          case _ => None
        }
        catalog.map(TransactionUtils.beginTransaction)
      } else {
        None
      }
    }
  }
  private def transactionOpt: Option[Transaction] = lazyTransactionOpt.get

  // Per-query analyzer: uses a transaction-aware CatalogManager when a transaction is active,
  // so that all catalog lookups and rule applications during analysis see the correct state
  // without relying on thread-local context. Any nested QueryExecution that is created during
  // analysis or execution of a transactional plan must receive this analyzer via analyzerOpt.
  private val lazyAnalyzer = LazyTry {
    analyzerOpt.getOrElse {
      transactionOpt match {
        case Some(txn) =>
          sparkSession.sessionState.analyzer.withCatalogManager(
            catalogManager.withTransaction(txn))
        case None =>
          sparkSession.sessionState.analyzer
      }
    }
  }
  private def analyzer: Analyzer = lazyAnalyzer.get

  def assertAnalyzed(): Unit = {
    try {
      analyzed
    } catch {
      case e: AnalysisException =>
        // Because we do eager analysis for Dataframe, there will be no execution created after
        // AnalysisException occurs. So we need to explicitly create a new execution to post
        // start/end events to notify the listener and UI components.
        SQLExecution.withNewExecutionIdOnError(this, Some("analyze"))(e)
    }
  }

  def assertSupported(): Unit = withAbortTransactionOnFailure {
    if (sparkSession.sessionState.conf.isUnsupportedOperationCheckEnabled) {
      UnsupportedOperationChecker.checkForBatch(analyzed)
    }
  }

  /**
   * Execute the SQL script if the logical plan is a SQL script.
   * There are multiple cases, and they are originating from:
   *   - SparkSession.sql() - Spark and Spark Connect case
   *   - QueryRuntimePredictionUtils.getParsedPlanWithTracking() - DBSQL case
   */
  private val lazySqlScriptExecuted = LazyTry {
    logical match {
      case NameParameterizedQuery(compoundBody: CompoundBody, argNames, argValues) =>
        SqlScriptingExecution.executeSqlScript(
          session = sparkSession,
          script = compoundBody,
          args = argNames.zip(argValues).toMap)
      case compoundBody: CompoundBody =>
        SqlScriptingExecution.executeSqlScript(
          session = sparkSession,
          script = compoundBody)
      case _ => logical
    }
  }

  private def sqlScriptExecuted: LogicalPlan = lazySqlScriptExecuted.get

  private def assertSqlScriptExecuted(): Unit = sqlScriptExecuted

  private val lazyAnalyzed = LazyTry {
    // Execute the SQL script. Script doesn't need to go through the analyzer as Spark
    //   will run each statement as individual query.
    assertSqlScriptExecuted()

    try {
      val plan = executePhase(QueryPlanningTracker.ANALYSIS) {
        // We can't clone `logical` here, which will reset the `_analyzed` flag.
        analyzer.executeAndCheck(sqlScriptExecuted, tracker)
      }
      tracker.setAnalyzed(plan)
      plan
    } catch {
      case NonFatal(e) =>
        tracker.setAnalysisFailed(sqlScriptExecuted)
        throw e
    }
  }

  def analyzed: LogicalPlan = withAbortTransactionOnFailure {
    lazyAnalyzed.get
  }

  private val lazyCommandExecuted = LazyTry {
    mode match {
      case CommandExecutionMode.NON_ROOT => analyzed.mapChildren(eagerlyExecuteCommands)
      case CommandExecutionMode.ALL => eagerlyExecuteCommands(analyzed)
      case CommandExecutionMode.SKIP => analyzed
    }
  }

  def commandExecuted: LogicalPlan = withAbortTransactionOnFailure {
    lazyCommandExecuted.get
  }

  private def commandExecutionName(command: Command): String = command match {
    case _: CreateTableAsSelect => "create"
    case _: ReplaceTableAsSelect => "replace"
    case _: AppendData => "append"
    case _: OverwriteByExpression => "overwrite"
    case _: OverwritePartitionsDynamic => "overwritePartitions"
    case _ => "command"
  }

  private def eagerlyExecuteCommands(p: LogicalPlan) = {
    def eagerlyExecute(
        p: LogicalPlan,
        name: String,
        mode: CommandExecutionMode.Value): LogicalPlan = {
      // Since Command execution will eagerly take place here,
      // and in most cases be the bulk of time and effort,
      // with the rest of processing of the root plan being just outputting command results,
      // for eagerly executed commands we mark this place as beginning of execution.
      tracker.setReadyForExecution()
      val (qe, result) = QueryExecution.runCommand(
        sparkSession, p, name, refreshPhaseEnabled, mode, Some(shuffleCleanupMode), Some(analyzer))
      CommandResult(
        qe.analyzed.output,
        qe.commandExecuted,
        qe.executedPlan,
        result.toImmutableArraySeq)
    }
    p transformDown {
      case u @ Union(children, _, _) if children.forall(_.isInstanceOf[Command]) =>
        eagerlyExecute(u, "multi-commands", CommandExecutionMode.SKIP)
      case w @ WithCTE(u @ Union(children, _, _), _) if children.forall(_.isInstanceOf[Command]) =>
        eagerlyExecute(w, "multi-commands", CommandExecutionMode.SKIP)
      case c: Command =>
        val name = commandExecutionName(c)
        eagerlyExecute(c, name, CommandExecutionMode.NON_ROOT)
      case w @ WithCTE(c: Command, _) =>
        val name = commandExecutionName(c)
        eagerlyExecute(w, name, CommandExecutionMode.SKIP)
    }
  }

  // there may be delay between analysis and subsequent phases
  // therefore, refresh captured table versions to reflect latest data
  private val lazyTableVersionsRefreshed = LazyTry {
    if (refreshPhaseEnabled) {
      V2TableRefreshUtil.refresh(sparkSession, commandExecuted, versionedOnly = true)
    } else {
      commandExecuted
    }
  }

  private[sql] def tableVersionsRefreshed: LogicalPlan = lazyTableVersionsRefreshed.get

  private val lazyNormalized = LazyTry {
    QueryExecution.normalize(sparkSession, tableVersionsRefreshed, Some(tracker))
  }

  // The plan that has been normalized by custom rules, so that it's more likely to hit cache.
  def normalized: LogicalPlan = withAbortTransactionOnFailure {
    lazyNormalized.get
  }

  private val lazyWithCachedData = LazyTry {
    sparkSession.withActive {
      assertAnalyzed()
      assertSupported()

      // Clone the plan to avoid sharing the plan instance between different stages like
      // analyzing, optimizing and planning.
      val planToRewrite = normalized.clone()
      val cacheManager = sparkSession.sharedState.cacheManager
      transactionOpt match {
        case Some(txn) => cacheManager.useCachedData(planToRewrite, txn)
        case None => cacheManager.useCachedData(planToRewrite)
      }
    }
  }

  def withCachedData: LogicalPlan = withAbortTransactionOnFailure {
    lazyWithCachedData.get
  }

  def assertCommandExecuted(): Unit = commandExecuted

  private val lazyOptimizedPlan = LazyTry {
    // We need to materialize the commandExecuted here because optimizedPlan is also tracked under
    // the optimizing phase
    assertCommandExecuted()
    executePhase(QueryPlanningTracker.OPTIMIZATION) {
      // clone the plan to avoid sharing the plan instance between different stages like analyzing,
      // optimizing and planning.
      val plan =
        sparkSession.sessionState.optimizer.executeAndTrack(withCachedData.clone(), tracker)
      // We do not want optimized plans to be re-analyzed as literals that have been constant
      // folded and such can cause issues during analysis. While `clone` should maintain the
      // `analyzed` state of the LogicalPlan, we set the plan as analyzed here as well out of
      // paranoia.
      plan.setAnalyzed()
      plan
    }
  }

  def optimizedPlan: LogicalPlan = withAbortTransactionOnFailure {
    lazyOptimizedPlan.get
  }

  def assertOptimized(): Unit = optimizedPlan

  private val lazySparkPlan = LazyTry {
    // We need to materialize the optimizedPlan here because sparkPlan is also tracked under
    // the planning phase
    assertOptimized()
    val plan = executePhase(QueryPlanningTracker.PLANNING) {
      // Clone the logical plan here, in case the planner rules change the states of the logical
      // plan.
      QueryExecution.createSparkPlan(planner, optimizedPlan.clone())
    }
    attachTransaction(plan)
  }

  def sparkPlan: SparkPlan = withAbortTransactionOnFailure {
    lazySparkPlan.get
  }

  def assertSparkPlanPrepared(): Unit = sparkPlan

  private val lazyExecutedPlan = LazyTry {
    // We need to materialize the optimizedPlan here, before tracking the planning phase, to ensure
    // that the optimization time is not counted as part of the planning phase.
    assertOptimized()
    val plan = executePhase(QueryPlanningTracker.PLANNING) {
      // clone the plan to avoid sharing the plan instance between different stages like analyzing,
      // optimizing and planning.
      QueryExecution.prepareForExecution(preparations, sparkPlan.clone())
    }
    // Note: For eagerly executed command it might have already been called in
    // `eagerlyExecutedCommand` and is a noop here.
    tracker.setReadyForExecution()
    plan
  }

  // executedPlan should not be used to initialize any SparkPlan. It should be
  // only used for execution.
  def executedPlan: SparkPlan = withAbortTransactionOnFailure {
    lazyExecutedPlan.get
  }

  def assertExecutedPlanPrepared(): Unit = executedPlan

  val lazyToRdd = LazyTry {
    new SQLExecutionRDD(executedPlan.execute(), sparkSession.sessionState.conf)
  }

  /**
   * Internal version of the RDD. Avoids copies and has no schema.
   * Note for callers: Spark may apply various optimization including reusing object: this means
   * the row is valid only for the iteration it is retrieved. You should avoid storing row and
   * accessing after iteration. (Calling `collect()` is one of known bad usage.)
   * If you want to store these rows into collection, please apply some converter or copy row
   * which produces new object per iteration.
   * Given QueryExecution is not a public class, end users are discouraged to use this: please
   * use `Dataset.rdd` instead where conversion will be applied.
   */
  def toRdd: RDD[InternalRow] = withAbortTransactionOnFailure {
    lazyToRdd.get
  }

  private val observedMetricsLock = new Object

  /** Get the metrics observed during the execution of the query plan. */
  @GuardedBy("observedMetricsLock")
  def observedMetrics: Map[String, Row] = observedMetricsLock.synchronized {
    CollectMetricsExec.collect(executedPlan)
  }

  protected def preparations: Seq[Rule[SparkPlan]] = {
    QueryExecution.preparations(sparkSession,
      Option(InsertAdaptiveSparkPlan(AdaptiveExecutionContext(sparkSession, this))), false)
  }

  protected def executePhase[T](phase: String)(block: => T): T = sparkSession.withActive {
    QueryExecution.withInternalError(s"The Spark SQL phase $phase failed with an internal error.") {
      withQueryExecutionId(sparkSession) {
        tracker.measurePhase(phase)(block)
      }
    }
  }

  /**
   * Set the query execution id in thread-local properties while
   * executing the block. This is used by
   * [[org.apache.spark.sql.execution.metric.SQLLastAttemptAccumulator]] to associate
   * driver-side metric updates with a specific QueryExecution.
   */
  private[sql] def withQueryExecutionId[T](
      session: SparkSession)(block: => T): T = {
    val sc = session.sparkContext
    val oldId = sc.getLocalProperty(
      SparkContext.DATASET_QUERY_EXECUTION_ID_KEY)
    sc.setLocalProperty(
      SparkContext.DATASET_QUERY_EXECUTION_ID_KEY, id.toString)
    try {
      block
    } finally {
      sc.setLocalProperty(
        SparkContext.DATASET_QUERY_EXECUTION_ID_KEY, oldId)
    }
  }

  def simpleString: String = {
    val concat = new PlanStringConcat()
    simpleString(false, SQLConf.get.maxToStringFields, concat.append)
    withRedaction {
      concat.toString
    }
  }

  private def simpleString(
      formatted: Boolean,
      maxFields: Int,
      append: String => Unit): Unit = {
    append("== Physical Plan ==\n")
    if (formatted) {
      try {
        ExplainUtils.processPlan(executedPlan, append)
      } catch {
        case e: AnalysisException => append(e.toString)
        case e: IllegalArgumentException => append(e.toString)
      }
    } else {
      QueryPlan.append(executedPlan,
        append, verbose = false, addSuffix = false, maxFields = maxFields)
    }
    extendedExplainInfo(append, executedPlan)
    append("\n")
  }

  def explainString(mode: ExplainMode): String = {
    val concat = new PlanStringConcat()
    explainString(mode, SQLConf.get.maxToStringFields, concat.append)
    withRedaction {
      concat.toString
    }
  }

  /**
   * Returns the QueryExecution to use when generating an explain string.
   * Overridden by IncrementalExecution to reuse `this` so that the already-open transaction and
   * cached executedPlan are not duplicated.
   */
  protected def queryExecutionForExplain: QueryExecution = if (logical.isStreaming) {
    // This is used only by explaining `Dataset/DataFrame` created by `spark.readStream`, so the
    // output mode does not matter since there is no `Sink`.
    new IncrementalExecution(
      sparkSession, logical, OutputMode.Append(), "<unknown>",
      UUID.randomUUID, UUID.randomUUID, 0, None, OffsetSeqMetadata(0, 0),
      WatermarkPropagator.noop(), false, mode = this.mode)
  } else {
    this
  }

  private def explainString(mode: ExplainMode, maxFields: Int, append: String => Unit): Unit = {
    val queryExecution = queryExecutionForExplain

    mode match {
      case SimpleMode =>
        queryExecution.simpleString(false, maxFields, append)
      case ExtendedMode =>
        queryExecution.toString(maxFields, append)
      case CodegenMode =>
        try {
          org.apache.spark.sql.execution.debug.writeCodegen(append, queryExecution.executedPlan)
        } catch {
          case e: AnalysisException => append(e.toString)
        }
      case CostMode =>
        queryExecution.stringWithStats(maxFields, append)
      case FormattedMode =>
        queryExecution.simpleString(formatted = true, maxFields = maxFields, append)
    }
  }

  private def writePlans(append: String => Unit, maxFields: Int): Unit = {
    val (verbose, addSuffix) = (true, false)
    append("== Parsed Logical Plan ==\n")
    QueryPlan.append(logical, append, verbose, addSuffix, maxFields)
    append("\n== Analyzed Logical Plan ==\n")
    try {
      if (analyzed.output.nonEmpty) {
        append(
          truncatedString(
            analyzed.output.map(o => s"${o.name}: ${o.dataType.simpleString}"), ", ", maxFields)
        )
        append("\n")
      }
      QueryPlan.append(analyzed, append, verbose, addSuffix, maxFields)
      append("\n== Optimized Logical Plan ==\n")
      QueryPlan.append(optimizedPlan, append, verbose, addSuffix, maxFields)
      append("\n== Physical Plan ==\n")
      QueryPlan.append(executedPlan, append, verbose, addSuffix, maxFields)
      extendedExplainInfo(append, executedPlan)
    } catch {
      case e: AnalysisException => append(e.toString)
    }
  }

  override def toString: String = withRedaction {
    val concat = new PlanStringConcat()
    toString(SQLConf.get.maxToStringFields, concat.append)
    withRedaction {
      concat.toString
    }
  }

  private def toString(maxFields: Int, append: String => Unit): Unit = {
    writePlans(append, maxFields)
  }

  def stringWithStats: String = {
    val concat = new PlanStringConcat()
    stringWithStats(SQLConf.get.maxToStringFields, concat.append)
    withRedaction {
      concat.toString
    }
  }

  private def stringWithStats(maxFields: Int, append: String => Unit): Unit = {
    // trigger to compute stats for logical plans
    try {
      // This will trigger to compute stats for all the nodes in the plan, including subqueries,
      // if the stats doesn't exist in the statsCache and update the statsCache corresponding
      // to the node.
      optimizedPlan.collectWithSubqueries {
        case plan => plan.stats
      }
    } catch {
      case e: AnalysisException => append(e.toString + "\n")
    }
    // only show optimized logical plan and physical plan
    append("== Optimized Logical Plan ==\n")
    QueryPlan.append(optimizedPlan, append, verbose = true, addSuffix = true, maxFields)
    append("\n== Physical Plan ==\n")
    QueryPlan.append(executedPlan, append, verbose = true, addSuffix = false, maxFields)
    append("\n")
  }

  /**
   * Redact the sensitive information in the given string.
   */
  private def withRedaction(message: String): String = {
    Utils.redact(sparkSession.sessionState.conf.stringRedactionPattern, message)
  }

  /**
   * Determine the shuffle cleanup mode, based on the following order:
   * 1. input arg when constructing this QueryExecution
   * 2. the cleanup mode of the root execution
   * 3. SQLConf from SparkSession
   */
  lazy val shuffleCleanupMode: ShuffleCleanupMode =
    shuffleCleanupModeOpt.getOrElse(
      getShuffleCleanupModeFromRootExecution.getOrElse(
        QueryExecution.determineShuffleCleanupMode(sparkSession.sessionState.conf)))

  private def getShuffleCleanupModeFromRootExecution: Option[ShuffleCleanupMode] = {
    val rootExecutionIdStr = sparkSession.sparkContext.getLocalProperty(EXECUTION_ROOT_ID_KEY)
    if (rootExecutionIdStr != null) {
      val rootExecutionId = rootExecutionIdStr.toLong
      val rootExecution = SQLExecution.getQueryExecution(rootExecutionId)
      if (rootExecution != null) {
        rootExecution.shuffleCleanupModeOpt
      } else {
        None
      }
    } else {
      None
    }
  }

  def extendedExplainInfo(append: String => Unit, plan: SparkPlan): Unit = {
    val generators = sparkSession.sessionState.conf.getConf(SQLConf.EXTENDED_EXPLAIN_PROVIDERS)
      .getOrElse(Seq.empty)
    val extensions = Utils.loadExtensions(classOf[ExtendedExplainGenerator],
      generators,
      sparkSession.sparkContext.conf)
    if (extensions.nonEmpty) {
      extensions.foreach(extension =>
        try {
          append(s"\n== Extended Information (${extension.title}) ==\n")
          append(extension.generateExtendedInfo(plan))
        } catch {
          case NonFatal(e) => logWarning(log"Cannot use " +
            log"${MDC(EXTENDED_EXPLAIN_GENERATOR, extension)} to get extended information.", e)
        })
    }
  }

  /**
   * Runs the given block, aborting the active transaction if an exception is thrown.
   * If no transaction is active, the block is executed as-is.
   */
  private def withAbortTransactionOnFailure[T](block: => T): T = transactionOpt match {
    case Some(transaction) =>
      try block
      catch { case e: Throwable => TransactionUtils.abort(transaction); throw e }
    case None => block
  }


  /** Attaches a transaction to the given SparkPlan to the transactional execution nodes. */
  private def attachTransaction(plan: SparkPlan): SparkPlan = transactionOpt match {
    case Some(txn) => plan.transformDown {
      case w: TransactionalExec => w.withTransaction(Some(txn))
    }
    case None => plan
  }

  /** A special namespace for commands that can be used to debug query execution. */
  // scalastyle:off
  object debug {
  // scalastyle:on

    /**
     * Prints to stdout all the generated code found in this plan (i.e. the output of each
     * WholeStageCodegen subtree).
     */
    def codegen(): Unit = {
      // scalastyle:off println
      println(org.apache.spark.sql.execution.debug.codegenString(executedPlan))
      // scalastyle:on println
    }

    /**
     * Get WholeStageCodegenExec subtrees and the codegen in a query plan
     *
     * @return Sequence of WholeStageCodegen subtrees and corresponding codegen
     */
    def codegenToSeq(): Seq[(String, String, ByteCodeStats)] = {
      org.apache.spark.sql.execution.debug.codegenStringSeq(executedPlan)
    }

    /**
     * Dumps debug information about query execution into the specified file.
     *
     * @param path path of the file the debug info is written to.
     * @param maxFields maximum number of fields converted to string representation.
     * @param explainMode the explain mode to be used to generate the string
     *                    representation of the plan.
     */
    def toFile(
        path: String,
        maxFields: Int = Int.MaxValue,
        explainMode: Option[String] = None): Unit = {
      val filePath = new Path(path)
      val fs = filePath.getFileSystem(sparkSession.sessionState.newHadoopConf())
      val writer = new BufferedWriter(new OutputStreamWriter(fs.create(filePath)))
      try {
        val mode = explainMode.map(ExplainMode.fromString(_)).getOrElse(ExtendedMode)
        explainString(mode, maxFields, writer.write)
        if (mode != CodegenMode) {
          writer.write("\n== Whole Stage Codegen ==\n")
          org.apache.spark.sql.execution.debug.writeCodegen(writer.write, executedPlan)
        }
        log.info(s"Debug information was written at: $filePath")
      } finally {
        writer.close()
      }
    }
  }
}

/**
 * SPARK-35378: Commands should be executed eagerly so that something like `sql("INSERT ...")`
 * can trigger the table insertion immediately without a `.collect()`. To avoid end-less recursion
 * we should use `NON_ROOT` when recursively executing commands. Note that we can't execute
 * a query plan with leaf command nodes, because many commands return `GenericInternalRow`
 * and can't be put in a query plan directly, otherwise the query engine may cast
 * `GenericInternalRow` to `UnsafeRow` and fail. When running EXPLAIN, or commands inside other
 * command, we should use `SKIP` to not eagerly trigger the command execution.
 */
object CommandExecutionMode extends Enumeration {
  val SKIP, NON_ROOT, ALL = Value
}

/**
 * Modes for shuffle dependency cleanup.
 *
 * DoNotCleanup: Do not perform any cleanup.
 * SkipMigration: Shuffle dependencies will not be migrated at node decommissions.
 * RemoveShuffleFiles: Shuffle dependency files are removed at the end of SQL executions.
 */
sealed trait ShuffleCleanupMode

case object DoNotCleanup extends ShuffleCleanupMode

case object SkipMigration extends ShuffleCleanupMode

case object RemoveShuffleFiles extends ShuffleCleanupMode


object QueryExecution {
  private val _nextExecutionId = new AtomicLong(0)

  private def nextExecutionId: Long = _nextExecutionId.getAndIncrement

  private[execution] def create(
      sparkSession: SparkSession,
      logical: LogicalPlan,
      refreshPhaseEnabled: Boolean = true): QueryExecution = {
    new QueryExecution(
      sparkSession,
      logical,
      mode = CommandExecutionMode.ALL,
      shuffleCleanupModeOpt = Some(determineShuffleCleanupMode(sparkSession.sessionState.conf)),
      refreshPhaseEnabled = refreshPhaseEnabled)
  }

  /**
   * Construct a sequence of rules that are used to prepare a planned [[SparkPlan]] for execution.
   * These rules will make sure subqueries are planned, make sure the data partitioning and ordering
   * are correct, insert whole stage code gen, and try to reduce the work done by reusing exchanges
   * and subqueries.
   */
  private[execution] def preparations(
      sparkSession: SparkSession,
      adaptiveExecutionRule: Option[InsertAdaptiveSparkPlan] = None,
      subquery: Boolean): Seq[Rule[SparkPlan]] = {
    // `AdaptiveSparkPlanExec` is a leaf node. If inserted, all the following rules will be no-op
    // as the original plan is hidden behind `AdaptiveSparkPlanExec`.
    adaptiveExecutionRule.toSeq ++
    Seq(
      CoalesceBucketsInJoin,
      PlanDynamicPruningFilters(sparkSession),
      PlanSubqueries(sparkSession),
      RemoveRedundantProjects,
      EnsureRequirements(),
      // This rule must be run after `EnsureRequirements`.
      InsertSortForLimitAndOffset,
      // `ReplaceHashWithSortAgg` needs to be added after `EnsureRequirements` to guarantee the
      // sort order of each node is checked to be valid.
      ReplaceHashWithSortAgg,
      // `RemoveRedundantSorts` and `RemoveRedundantWindowGroupLimits` needs to be added after
      // `EnsureRequirements` to guarantee the same number of partitions when instantiating
      // PartitioningCollection.
      RemoveRedundantSorts,
      RemoveRedundantWindowGroupLimits,
      DisableUnnecessaryBucketedScan,
      ApplyColumnarRulesAndInsertTransitions(
        sparkSession.sessionState.columnarRules, outputsColumnar = false),
      CollapseCodegenStages()) ++
      (if (subquery) {
        Nil
      } else {
        Seq(ReuseExchangeAndSubquery)
      })
  }

  /**
   * Prepares a planned [[SparkPlan]] for execution by inserting shuffle operations and internal
   * row format conversions as needed.
   */
  private[execution] def prepareForExecution(
      preparations: Seq[Rule[SparkPlan]],
      plan: SparkPlan): SparkPlan = {
    val planChangeLogger = new PlanChangeLogger[SparkPlan]()
    val preparedPlan = preparations.foldLeft(plan) { case (sp, rule) =>
      val result = rule.apply(sp)
      planChangeLogger.logRule(rule.ruleName, sp, result)
      result
    }
    planChangeLogger.logBatch("Preparations", plan, preparedPlan)
    preparedPlan
  }

  /**
   * Transform a [[LogicalPlan]] into a [[SparkPlan]].
   *
   * Note that the returned physical plan still needs to be prepared for execution.
   */
  def createSparkPlan(
      planner: SparkPlanner,
      plan: LogicalPlan): SparkPlan = {
    // TODO: We use next(), i.e. take the first plan returned by the planner, here for now,
    //       but we will implement to choose the best plan.
    planner.plan(ReturnAnswer(plan)).next()
  }

  /**
   * Prepare the [[SparkPlan]] for execution.
   */
  def prepareExecutedPlan(spark: SparkSession, plan: SparkPlan): SparkPlan = {
    prepareForExecution(preparations(spark, subquery = true), plan)
  }

  /**
   * Transform the subquery's [[LogicalPlan]] into a [[SparkPlan]] and prepare the resulting
   * [[SparkPlan]] for execution.
   */
  def prepareExecutedPlan(spark: SparkSession, plan: LogicalPlan): SparkPlan = {
    val sparkPlan = createSparkPlan(spark.sessionState.planner, plan.clone())
    prepareExecutedPlan(spark, sparkPlan)
  }

  /**
   * Prepare the [[SparkPlan]] for execution using exists adaptive execution context.
   * This method is only called by [[PlanAdaptiveDynamicPruningFilters]].
   */
  def prepareExecutedPlan(
      plan: LogicalPlan,
      context: AdaptiveExecutionContext): SparkPlan = {
    val sparkPlan = createSparkPlan(context.session.sessionState.planner, plan.clone())
    val preparationRules =
      preparations(context.session, Option(InsertAdaptiveSparkPlan(context)), true)
    prepareForExecution(preparationRules, sparkPlan.clone())
  }

  /**
   * Marks null pointer exceptions, asserts and scala match errors as internal errors
   */
  private[sql] def isInternalError(e: Throwable): Boolean = e match {
    case _: java.lang.NullPointerException => true
    case _: java.lang.AssertionError => true
    case _: scala.MatchError => true
    case _ => false
  }

  /**
   * Converts marked exceptions from [[isInternalError]] to internal errors.
   */
  private[sql] def toInternalError(msg: String, e: Throwable): Throwable = {
    if (isInternalError(e)) {
      SparkException.internalError(
        msg + " You hit a bug in Spark or the Spark plugins you use. Please, report this bug " +
          "to the corresponding communities or vendors, and provide the full stack trace.",
        e)
    } else {
      e
    }
  }

  /**
   * Catches marked exceptions from [[isInternalError]], and converts them to internal errors.
   */
  private[sql] def withInternalError[T](msg: String)(block: => T): T = {
    try {
      block
    } catch {
      case e: Throwable => throw toInternalError(msg, e)
    }
  }

  def normalize(
      session: SparkSession,
      plan: LogicalPlan,
      tracker: Option[QueryPlanningTracker] = None): LogicalPlan = {
    val normalizationRules = session.sessionState.planNormalizationRules
    if (normalizationRules.isEmpty) {
      plan
    } else {
      val planChangeLogger = new PlanChangeLogger[LogicalPlan]()
      val normalized = normalizationRules.foldLeft(plan) { (p, rule) =>
        val startTime = System.nanoTime()
        val result = rule.apply(p)
        val runTime = System.nanoTime() - startTime
        val effective = !result.fastEquals(p)
        tracker.foreach(_.recordRuleInvocation(rule.ruleName, runTime, effective))
        planChangeLogger.logRule(rule.ruleName, p, result)
        result
      }
      planChangeLogger.logBatch("Plan Normalization", plan, normalized)
      normalized
    }
  }

  def determineShuffleCleanupMode(conf: SQLConf): ShuffleCleanupMode = {
    if (conf.getConf(SQLConf.CLASSIC_SHUFFLE_DEPENDENCY_FILE_CLEANUP_ENABLED)) {
      RemoveShuffleFiles
    } else {
      DoNotCleanup
    }
  }

  /**
   * Determines whether the given unresolved plan is a SQL script.
   * @param plan Logical plan to check.
   * @return True if the plan is a SQL script, False otherwise.
   */
  def isUnresolvedPlanSqlScript(plan: LogicalPlan): Boolean = {
    plan match {
      case _: CompoundBody => true
      case NameParameterizedQuery(_: CompoundBody, _, _) => true
      case _ => false
    }
  }

  def runCommand(
      sparkSession: SparkSession,
      command: LogicalPlan,
      name: String,
      refreshPhaseEnabled: Boolean = true,
      mode: CommandExecutionMode.Value = CommandExecutionMode.SKIP,
      shuffleCleanupModeOpt: Option[ShuffleCleanupMode] = None,
      analyzerOpt: Option[Analyzer] = None)
    : (QueryExecution, Array[InternalRow]) = {
    val qe = new QueryExecution(
      sparkSession,
      command,
      mode = mode,
      shuffleCleanupModeOpt = shuffleCleanupModeOpt,
      refreshPhaseEnabled = refreshPhaseEnabled,
      analyzerOpt = analyzerOpt)
    val result = QueryExecution.withInternalError(s"Executed $name failed.") {
      SQLExecution.withNewExecutionId(qe, Some(name)) {
        qe.executedPlan.executeCollect()
      }
    }
    (qe, result)
  }
}
