[SPARK-41631][SQL] Support implicit lateral column alias resolution on Aggregate

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

@@ -1744,7 +1744,7 @@ class Analyzer(override val catalogManager: CatalogManager)
      *                    the current plan is needed to first try resolving the attribute by its
      *                    children
      */
-    private def wrapLCARefHelper(
+    private def wrapLCARef(
         e: NamedExpression,
         currentPlan: LogicalPlan,
         aliasMap: CaseInsensitiveMap[Seq[AliasEntry]]): NamedExpression = {
@@ -1763,9 +1763,14 @@ class Analyzer(override val catalogManager: CatalogManager)
             case _ => u
           }
         case o: OuterReference
-          if aliasMap.contains(o.nameParts.map(_.head).getOrElse(o.name)) =>
+          if aliasMap.contains(
+            o.getTagValue(ResolveLateralColumnAliasReference.NAME_PARTS_FROM_UNRESOLVED_ATTR)
+              .map(_.head)
+              .getOrElse(o.name)) =>
           // handle OuterReference exactly same as UnresolvedAttribute
-          val nameParts = o.nameParts.getOrElse(Seq(o.name))
+          val nameParts = o
+            .getTagValue(ResolveLateralColumnAliasReference.NAME_PARTS_FROM_UNRESOLVED_ATTR)
+            .getOrElse(Seq(o.name))
           val aliases = aliasMap.get(nameParts.head).get
           aliases.size match {
             case n if n > 1 =>
@@ -1789,15 +1794,15 @@ class Analyzer(override val catalogManager: CatalogManager)
             var aliasMap = CaseInsensitiveMap(Map[String, Seq[AliasEntry]]())
             val newProjectList = projectList.zipWithIndex.map {
               case (a: Alias, idx) =>
-                val lcaWrapped = wrapLCARefHelper(a, p, aliasMap).asInstanceOf[Alias]
+                val lcaWrapped = wrapLCARef(a, p, aliasMap).asInstanceOf[Alias]
                 // Insert the LCA-resolved alias instead of the unresolved one into map. If it is
                 // resolved, it can be referenced as LCA by later expressions (chaining).
                 // Unresolved Alias is also added to the map to perform ambiguous name check, but
                 // only resolved alias can be LCA.
                 aliasMap = insertIntoAliasMap(lcaWrapped, idx, aliasMap)
                 lcaWrapped
               case (e, _) =>
-                wrapLCARefHelper(e, p, aliasMap)
+                wrapLCARef(e, p, aliasMap)
             }
             p.copy(projectList = newProjectList)
 

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

@@ -644,8 +644,10 @@ trait CheckAnalysis extends PredicateHelper with LookupCatalog with QueryErrorsB
             projectList.foreach(_.transformDownWithPruning(
               _.containsPattern(LATERAL_COLUMN_ALIAS_REFERENCE)) {
               case lcaRef: LateralColumnAliasReference if p.resolved =>
-                failUnresolvedAttribute(
-                  p, UnresolvedAttribute(lcaRef.nameParts), "UNRESOLVED_COLUMN")
+                throw SparkException.internalError("Resolved Project should not contain " +
+                  s"any LateralColumnAliasReference.\nDebugging information: plan: $p",
+                  context = lcaRef.origin.getQueryContext,
+                  summary = lcaRef.origin.context.summary)
             })
 
           case j: Join if !j.duplicateResolved =>
@@ -730,8 +732,10 @@ trait CheckAnalysis extends PredicateHelper with LookupCatalog with QueryErrorsB
             aggList.foreach(_.transformDownWithPruning(
               _.containsPattern(LATERAL_COLUMN_ALIAS_REFERENCE)) {
               case lcaRef: LateralColumnAliasReference =>
-                failUnresolvedAttribute(
-                  agg, UnresolvedAttribute(lcaRef.nameParts), "UNRESOLVED_COLUMN")
+                throw SparkException.internalError("Resolved Aggregate should not contain " +
+                  s"any LateralColumnAliasReference.\nDebugging information: plan: $agg",
+                  context = lcaRef.origin.getQueryContext,
+                  summary = lcaRef.origin.context.summary)
             })
 
           case _ => // Analysis successful!

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

@@ -21,6 +21,7 @@ import org.apache.spark.sql.catalyst.expressions.{Alias, AttributeMap, Expressio
 import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
 import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan, Project}
 import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
 import org.apache.spark.sql.catalyst.trees.TreePattern.LATERAL_COLUMN_ALIAS_REFERENCE
 import org.apache.spark.sql.errors.QueryCompilationErrors
 import org.apache.spark.sql.internal.SQLConf
@@ -93,6 +94,13 @@ import org.apache.spark.sql.internal.SQLConf
 object ResolveLateralColumnAliasReference extends Rule[LogicalPlan] {
   case class AliasEntry(alias: Alias, index: Int)
 
+  /**
+   * A tag to store the nameParts from the original unresolved attribute.
+   * It is set for [[OuterReference]], used in the current rule to convert [[OuterReference]] back
+   * to [[LateralColumnAliasReference]].
+   */
+  val NAME_PARTS_FROM_UNRESOLVED_ATTR = TreeNodeTag[Seq[String]]("name_parts_from_unresolved_attr")
+
   override def apply(plan: LogicalPlan): LogicalPlan = {
     if (!conf.getConf(SQLConf.LATERAL_COLUMN_ALIAS_IMPLICIT_ENABLED)) {
       plan

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala

@@ -424,18 +424,8 @@ case class OuterReference(e: NamedExpression)
   override def qualifier: Seq[String] = e.qualifier
   override def exprId: ExprId = e.exprId
   override def toAttribute: Attribute = e.toAttribute
-  override def newInstance(): NamedExpression =
-    OuterReference(e.newInstance()).setNameParts(nameParts)
+  override def newInstance(): NamedExpression = OuterReference(e.newInstance())
   final override val nodePatterns: Seq[TreePattern] = Seq(OUTER_REFERENCE)
-
-  // optional field, the original name parts of UnresolvedAttribute before it is resolved to
-  // OuterReference. Used in rule ResolveLateralColumnAlias to convert OuterReference back to
-  // LateralColumnAliasReference.
-  var nameParts: Option[Seq[String]] = None
-  def setNameParts(newNameParts: Option[Seq[String]]): OuterReference = {
-    nameParts = newNameParts
-    this
-  }
 }
 
 /**

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/subquery.scala

@@ -19,6 +19,7 @@ package org.apache.spark.sql.catalyst.expressions
 
 import scala.collection.mutable.ArrayBuffer
 
+import org.apache.spark.sql.catalyst.analysis.ResolveLateralColumnAliasReference.NAME_PARTS_FROM_UNRESOLVED_ATTR
 import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
 import org.apache.spark.sql.catalyst.plans.QueryPlan
 import org.apache.spark.sql.catalyst.plans.logical.{Filter, HintInfo, LogicalPlan}
@@ -159,7 +160,11 @@ object SubExprUtils extends PredicateHelper {
    * Wrap attributes in the expression with [[OuterReference]]s.
    */
   def wrapOuterReference[E <: Expression](e: E, nameParts: Option[Seq[String]] = None): E = {
-    e.transform { case a: Attribute => OuterReference(a).setNameParts(nameParts) }.asInstanceOf[E]
+    e.transform { case a: Attribute =>
+      val o = OuterReference(a)
+      nameParts.map(o.setTagValue(NAME_PARTS_FROM_UNRESOLVED_ATTR, _))
+      o
+    }.asInstanceOf[E]
   }
 
   /**