[SPARK-45959][SQL] Improving performance when addition of 1 column at a time causes increase in the LogicalPlan tree depth

connector/connect/common/src/test/resources/query-tests/explain-results/crosstab.explain

@@ -1,5 +1,4 @@
 Project [a_b#0]
-+- Project [a_b#0]
-   +- Aggregate [a_b#0], [a_b#0, pivotfirst(__pivot_col#0, count(1) AS count#0L, 0, 0) AS __pivot_count(1) AS count AS `count(1) AS count`#0]
-      +- Aggregate [CASE WHEN isnull(a#0) THEN null ELSE cast(a#0 as string) END, CASE WHEN isnull(b#0) THEN null ELSE regexp_replace(cast(b#0 as string), `, , 1) END], [CASE WHEN isnull(a#0) THEN null ELSE cast(a#0 as string) END AS a_b#0, CASE WHEN isnull(b#0) THEN null ELSE regexp_replace(cast(b#0 as string), `, , 1) END AS __pivot_col#0, count(1) AS count(1) AS count#0L]
-         +- LocalRelation <empty>, [id#0L, a#0, b#0]
++- Aggregate [a_b#0], [a_b#0, pivotfirst(__pivot_col#0, count(1) AS count#0L, 0, 0) AS __pivot_count(1) AS count AS `count(1) AS count`#0]
+   +- Aggregate [CASE WHEN isnull(a#0) THEN null ELSE cast(a#0 as string) END, CASE WHEN isnull(b#0) THEN null ELSE regexp_replace(cast(b#0 as string), `, , 1) END], [CASE WHEN isnull(a#0) THEN null ELSE cast(a#0 as string) END AS a_b#0, CASE WHEN isnull(b#0) THEN null ELSE regexp_replace(cast(b#0 as string), `, , 1) END AS __pivot_col#0, count(1) AS count(1) AS count#0L]
+      +- LocalRelation <empty>, [id#0L, a#0, b#0]

connector/connect/common/src/test/resources/query-tests/explain-results/describe.explain

@@ -1,6 +1,5 @@
 Project [summary#0, element_at(id#0, summary#0, None, false) AS id#0, element_at(b#0, summary#0, None, false) AS b#0]
-+- Project [id#0, b#0, summary#0]
-   +- Generate explode([count,mean,stddev,min,max]), false, [summary#0]
-      +- Aggregate [map(cast(count as string), cast(count(id#0L) as string), cast(mean as string), cast(avg(id#0L) as string), cast(stddev as string), cast(stddev(cast(id#0L as double)) as string), cast(min as string), cast(min(id#0L) as string), cast(max as string), cast(max(id#0L) as string)) AS id#0, map(cast(count as string), cast(count(b#0) as string), cast(mean as string), cast(avg(b#0) as string), cast(stddev as string), cast(stddev(b#0) as string), cast(min as string), cast(min(b#0) as string), cast(max as string), cast(max(b#0) as string)) AS b#0]
-         +- Project [id#0L, b#0]
-            +- LocalRelation <empty>, [id#0L, a#0, b#0]
++- Generate explode([count,mean,stddev,min,max]), false, [summary#0]
+   +- Aggregate [map(cast(count as string), cast(count(id#0L) as string), cast(mean as string), cast(avg(id#0L) as string), cast(stddev as string), cast(stddev(cast(id#0L as double)) as string), cast(min as string), cast(min(id#0L) as string), cast(max as string), cast(max(id#0L) as string)) AS id#0, map(cast(count as string), cast(count(b#0) as string), cast(mean as string), cast(avg(b#0) as string), cast(stddev as string), cast(stddev(b#0) as string), cast(min as string), cast(min(b#0) as string), cast(max as string), cast(max(b#0) as string)) AS b#0]
+      +- Project [id#0L, b#0]
+         +- LocalRelation <empty>, [id#0L, a#0, b#0]

connector/connect/common/src/test/resources/query-tests/explain-results/summary.explain

@@ -1,5 +1,4 @@
 Project [summary#0, element_at(id#0, summary#0, None, false) AS id#0, element_at(a#0, summary#0, None, false) AS a#0, element_at(b#0, summary#0, None, false) AS b#0]
-+- Project [id#0, a#0, b#0, summary#0]
-   +- Generate explode([mean,min]), false, [summary#0]
-      +- Aggregate [map(cast(mean as string), cast(avg(id#0L) as string), cast(min as string), cast(min(id#0L) as string)) AS id#0, map(cast(mean as string), cast(avg(a#0) as string), cast(min as string), cast(min(a#0) as string)) AS a#0, map(cast(mean as string), cast(avg(b#0) as string), cast(min as string), cast(min(b#0) as string)) AS b#0]
-         +- LocalRelation <empty>, [id#0L, a#0, b#0]
++- Generate explode([mean,min]), false, [summary#0]
+   +- Aggregate [map(cast(mean as string), cast(avg(id#0L) as string), cast(min as string), cast(min(id#0L) as string)) AS id#0, map(cast(mean as string), cast(avg(a#0) as string), cast(min as string), cast(min(a#0) as string)) AS a#0, map(cast(mean as string), cast(avg(b#0) as string), cast(min as string), cast(min(b#0) as string)) AS b#0]
+      +- LocalRelation <empty>, [id#0L, a#0, b#0]

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala

@@ -59,6 +59,7 @@ import org.apache.spark.sql.types._
 import org.apache.spark.sql.types.DayTimeIntervalType.DAY
 import org.apache.spark.sql.util.CaseInsensitiveStringMap
 import org.apache.spark.util.ArrayImplicits._
+import org.apache.spark.util.Utils
 
 /**
  * A trivial [[Analyzer]] with a dummy [[SessionCatalog]] and
@@ -212,7 +213,12 @@ class Analyzer(override val catalogManager: CatalogManager) extends RuleExecutor
     AnalysisHelper.markInAnalyzer {
       val analyzed = executeAndTrack(plan, tracker)
       checkAnalysis(analyzed)
-      analyzed
+      val excludedPostAnalysisRulesConf =
+        conf.postAnalysisExcludesRules.toSeq.flatMap(Utils.stringToSeq)
+      postAnalysisEarlyOptimizationRules.filterNot(
+        rule => excludedPostAnalysisRulesConf.contains(rule.ruleName)).foldLeft(analyzed) {
+        case(rs, rule) => rule(rs)
+      }
     }
   }
 
@@ -248,6 +254,8 @@ class Analyzer(override val catalogManager: CatalogManager) extends RuleExecutor
    */
   val postHocResolutionRules: Seq[Rule[LogicalPlan]] = Nil
 
+  val postAnalysisEarlyOptimizationRules: Seq[Rule[LogicalPlan]] = Nil
+
   private def typeCoercionRules(): List[Rule[LogicalPlan]] = if (conf.ansiEnabled) {
     AnsiTypeCoercion.typeCoercionRules
   } else {
@@ -356,6 +364,8 @@ class Analyzer(override val catalogManager: CatalogManager) extends RuleExecutor
     Batch("HandleSpecialCommand", Once,
       HandleSpecialCommand),
     Batch("Remove watermark for batch query", Once,
+      EliminateEventTimeWatermark),
+      Batch("Remove watermark for batch query", Once,
       EliminateEventTimeWatermark)
   )
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala

@@ -59,24 +59,39 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
           val (newExprs, newChild) = {
             // Resolving expressions against current plan.
             val maybeResolvedExprs = exprs.map(resolveExpressionByPlanOutput(_, u))
+
             // Recursively resolving expressions on the child of current plan.
             resolveExprsAndAddMissingAttrs(maybeResolvedExprs, u.child)
           }
           // If some attributes used by expressions are resolvable only on the rewritten child
           // plan, we need to add them into original projection.
-          lazy val missingAttrs =
-            (AttributeSet(newExprs) -- u.outputSet).intersect(newChild.outputSet)
+          val (missingAttrsFromOutput, missingAttrsFromDroppedAttr) = {
+            val missing1 = AttributeSet(newExprs) -- u.outputSet
+            val fulfilledFromOutput = missing1.intersect(newChild.outputSet)
+            val missing2 = missing1 -- fulfilledFromOutput
+            val fulfilledFromDroppedCol = missing2.intersect(u.getTagValue(
+              LogicalPlan.DROPPED_NAMED_EXPRESSIONS).
+              map(sq => AttributeSet(sq.map(_.toAttribute))).getOrElse(AttributeSet.empty))
+            fulfilledFromOutput -> fulfilledFromDroppedCol
+          }
           u match {
             case p: Project =>
-              val newProject = Project(p.projectList ++ missingAttrs, newChild)
+              val droppedNamedExprs = p.getTagValue(LogicalPlan.DROPPED_NAMED_EXPRESSIONS).
+                getOrElse(Seq.empty)
+              val newProject = Project(p.projectList ++ missingAttrsFromOutput ++
+                missingAttrsFromDroppedAttr.map(attr =>
+                  droppedNamedExprs.find(_.toAttribute.canonicalized == attr.canonicalized).get),
+                  newChild)
               newProject.copyTagsFrom(p)
               (newExprs, newProject)
 
-            case a @ Aggregate(groupExprs, aggExprs, child) =>
-              if (missingAttrs.forall(attr => groupExprs.exists(_.semanticEquals(attr)))) {
+            case a @ Aggregate(groupExprs, aggExprs, _) =>
+              if (missingAttrsFromOutput.forall(attr =>
+                groupExprs.exists(_.semanticEquals(attr)))) {
                 // All the missing attributes are grouping expressions, valid case.
                 (newExprs,
-                  a.copy(aggregateExpressions = aggExprs ++ missingAttrs, child = newChild))
+                  a.copy(aggregateExpressions = aggExprs ++ missingAttrsFromOutput,
+                    child = newChild))
               } else {
                 // Need to add non-grouping attributes, invalid case.
                 (exprs, a)

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala

@@ -136,6 +136,9 @@ abstract class LogicalPlan
 
   private[this] lazy val outputAttributes = AttributeSeq.fromNormalOutput(output)
 
+  private[this] lazy val droppedAttributes = this.getTagValue(
+    LogicalPlan.DROPPED_NAMED_EXPRESSIONS).map(_.map(_.toAttribute)).getOrElse(Seq.empty)
+
   private[this] lazy val outputMetadataAttributes = AttributeSeq(metadataOutput)
 
   /**
@@ -158,7 +161,8 @@ abstract class LogicalPlan
       nameParts: Seq[String],
       resolver: Resolver): Option[NamedExpression] =
     outputAttributes.resolve(nameParts, resolver)
-      .orElse(outputMetadataAttributes.resolve(nameParts, resolver))
+      .orElse(outputMetadataAttributes.resolve(nameParts, resolver)).orElse(
+      droppedAttributes.resolve(nameParts, resolver))
 
   /**
    * Given an attribute name, split it to name parts by dot, but
@@ -202,6 +206,8 @@ object LogicalPlan {
   //    to the old code path.
   private[spark] val PLAN_ID_TAG = TreeNodeTag[Long]("plan_id")
   private[spark] val IS_METADATA_COL = TreeNodeTag[Unit]("is_metadata_col")
+  private[spark] val DROPPED_NAMED_EXPRESSIONS =
+    TreeNodeTag[Seq[NamedExpression]]("dropped_namedexprs")
 }
 
 /**

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -238,6 +238,14 @@ object SQLConf {
     }
   }
 
+  val EXCLUDE_POST_ANALYSIS_RULES =
+    buildConf("spark.sql.analyzer.excludePostAnalysisRules")
+    .internal()
+    .doc("The names of the comma separated post analysis rules to be excluded")
+    .version("3.5.0")
+    .stringConf
+    .createOptional
+
   val ANALYZER_MAX_ITERATIONS = buildConf("spark.sql.analyzer.maxIterations")
     .internal()
     .doc("The max number of iterations the analyzer runs.")
@@ -5159,6 +5167,8 @@ class SQLConf extends Serializable with Logging with SqlApiConf {
 
   def analyzerMaxIterations: Int = getConf(ANALYZER_MAX_ITERATIONS)
 
+  def postAnalysisExcludesRules: Option[String] = getConf(EXCLUDE_POST_ANALYSIS_RULES)
+
   def optimizerExcludedRules: Option[String] = getConf(OPTIMIZER_EXCLUDED_RULES)
 
   def optimizerMaxIterations: Int = getConf(OPTIMIZER_MAX_ITERATIONS)

sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala

@@ -3890,7 +3890,8 @@ class Dataset[T] private[sql](
    */
   def storageLevel: StorageLevel = {
     sparkSession.sharedState.cacheManager.lookupCachedData(this).map { cachedData =>
-      cachedData.cachedRepresentation.cacheBuilder.storageLevel
+      cachedData.cachedRepresentation.fold(CacheManager.inMemoryRelationExtractor, identity).
+        cacheBuilder.storageLevel
     }.getOrElse(StorageLevel.NONE)
   }