Merge remote-tracking branch 'apache/master' into SPARK-27561-agg

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

@@ -1763,141 +1763,6 @@ class Analyzer(override val catalogManager: CatalogManager)
     }
   }
 
-  /**
-   * The first phase to resolve lateral column alias. See comments in
-   * [[ResolveLateralColumnAliasReference]] for more detailed explanation.
-   */
-  object WrapLateralColumnAliasReference extends Rule[LogicalPlan] {
-    import ResolveLateralColumnAliasReference.AliasEntry
-
-    private def insertIntoAliasMap(
-        a: Alias,
-        idx: Int,
-        aliasMap: CaseInsensitiveMap[Seq[AliasEntry]]): CaseInsensitiveMap[Seq[AliasEntry]] = {
-      val prevAliases = aliasMap.getOrElse(a.name, Seq.empty[AliasEntry])
-      aliasMap + (a.name -> (prevAliases :+ AliasEntry(a, idx)))
-    }
-
-    /**
-     * Use the given lateral alias to resolve the unresolved attribute with the name parts.
-     *
-     * Construct a dummy plan with the given lateral alias as project list, use the output of the
-     * plan to resolve.
-     * @return The resolved [[LateralColumnAliasReference]] if succeeds. None if fails to resolve.
-     */
-    private def resolveByLateralAlias(
-        nameParts: Seq[String], lateralAlias: Alias): Option[LateralColumnAliasReference] = {
-      val resolvedAttr = resolveExpressionByPlanOutput(
-        expr = UnresolvedAttribute(nameParts),
-        plan = LocalRelation(Seq(lateralAlias.toAttribute)),
-        throws = false
-      ).asInstanceOf[NamedExpression]
-      if (resolvedAttr.resolved) {
-        Some(LateralColumnAliasReference(resolvedAttr, nameParts, lateralAlias.toAttribute))
-      } else {
-        None
-      }
-    }
-
-    /**
-     * Recognize all the attributes in the given expression that reference lateral column aliases
-     * by looking up the alias map. Resolve these attributes and replace by wrapping with
-     * [[LateralColumnAliasReference]].
-     *
-     * @param currentPlan Because lateral alias has lower resolution priority than table columns,
-     *                    the current plan is needed to first try resolving the attribute by its
-     *                    children
-     */
-    private def wrapLCARef(
-        e: NamedExpression,
-        currentPlan: LogicalPlan,
-        aliasMap: CaseInsensitiveMap[Seq[AliasEntry]]): NamedExpression = {
-      e.transformWithPruning(_.containsAnyPattern(UNRESOLVED_ATTRIBUTE, OUTER_REFERENCE)) {
-        case u: UnresolvedAttribute if aliasMap.contains(u.nameParts.head) &&
-          resolveExpressionByPlanChildren(u, currentPlan).isInstanceOf[UnresolvedAttribute] =>
-          val aliases = aliasMap.get(u.nameParts.head).get
-          aliases.size match {
-            case n if n > 1 =>
-              throw QueryCompilationErrors.ambiguousLateralColumnAliasError(u.name, n)
-            case n if n == 1 && aliases.head.alias.resolved =>
-              // Only resolved alias can be the lateral column alias
-              // The lateral alias can be a struct and have nested field, need to construct
-              // a dummy plan to resolve the expression
-              resolveByLateralAlias(u.nameParts, aliases.head.alias).getOrElse(u)
-            case _ => u
-          }
-        case o: OuterReference
-          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
-            .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 =>
-              throw QueryCompilationErrors.ambiguousLateralColumnAliasError(nameParts, n)
-            case n if n == 1 && aliases.head.alias.resolved =>
-              resolveByLateralAlias(nameParts, aliases.head.alias).getOrElse(o)
-            case _ => o
-          }
-      }.asInstanceOf[NamedExpression]
-    }
-
-    override def apply(plan: LogicalPlan): LogicalPlan = {
-      if (!conf.getConf(SQLConf.LATERAL_COLUMN_ALIAS_IMPLICIT_ENABLED)) {
-        plan
-      } else {
-        plan.resolveOperatorsUpWithPruning(
-          _.containsAnyPattern(UNRESOLVED_ATTRIBUTE, OUTER_REFERENCE), ruleId) {
-          case p @ Project(projectList, _) if p.childrenResolved
-            && !ResolveReferences.containsStar(projectList)
-            && projectList.exists(_.containsAnyPattern(UNRESOLVED_ATTRIBUTE, OUTER_REFERENCE)) =>
-            var aliasMap = CaseInsensitiveMap(Map[String, Seq[AliasEntry]]())
-            val newProjectList = projectList.zipWithIndex.map {
-              case (a: Alias, idx) =>
-                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, _) =>
-                wrapLCARef(e, p, aliasMap)
-            }
-            p.copy(projectList = newProjectList)
-
-          // Implementation notes:
-          // In Aggregate, introducing and wrapping this resolved leaf expression
-          // LateralColumnAliasReference is especially needed because it needs an accurate condition
-          // to trigger adding a Project above and extracting and pushing down aggregate functions
-          // or grouping expressions. Such operation can only be done once. With this
-          // LateralColumnAliasReference, that condition can simply be when the whole Aggregate is
-          // resolved. Otherwise, it can't tell if all aggregate functions are created and
-          // resolved so that it can start the extraction, because the lateral alias reference is
-          // unresolved and can be the argument to functions, blocking the resolution of functions.
-          case agg @ Aggregate(_, aggExprs, _) if agg.childrenResolved
-              && !ResolveReferences.containsStar(aggExprs)
-              && aggExprs.exists(_.containsAnyPattern(UNRESOLVED_ATTRIBUTE, OUTER_REFERENCE)) =>
-
-            var aliasMap = CaseInsensitiveMap(Map[String, Seq[AliasEntry]]())
-            val newAggExprs = aggExprs.zipWithIndex.map {
-              case (a: Alias, idx) =>
-                val lcaWrapped = wrapLCARef(a, agg, aliasMap).asInstanceOf[Alias]
-                aliasMap = insertIntoAliasMap(lcaWrapped, idx, aliasMap)
-                lcaWrapped
-              case (e, _) =>
-                wrapLCARef(e, agg, aliasMap)
-            }
-            agg.copy(aggregateExpressions = newAggExprs)
-        }
-      }
-    }
-  }
-
   private def containsDeserializer(exprs: Seq[Expression]): Boolean = {
     exprs.exists(_.exists(_.isInstanceOf[UnresolvedDeserializer]))
   }

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

@@ -17,13 +17,13 @@
 
 package org.apache.spark.sql.catalyst.analysis
 
-import org.apache.spark.sql.catalyst.expressions.{Alias, AttributeMap, Expression, LateralColumnAliasReference, NamedExpression}
+import org.apache.spark.sql.catalyst.expressions.{Alias, AttributeMap, Expression, LateralColumnAliasReference, NamedExpression, OuterReference}
 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.plans.logical.{Aggregate, LocalRelation, 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.catalyst.util.toPrettySQL
+import org.apache.spark.sql.catalyst.trees.TreePattern.{LATERAL_COLUMN_ALIAS_REFERENCE, OUTER_REFERENCE, UNRESOLVED_ATTRIBUTE}
+import org.apache.spark.sql.catalyst.util.{toPrettySQL, CaseInsensitiveMap}
 import org.apache.spark.sql.errors.QueryCompilationErrors
 import org.apache.spark.sql.internal.SQLConf
 
@@ -144,7 +144,7 @@ object WrapLateralColumnAliasReference extends Rule[LogicalPlan] {
         val aliases = aliasMap.get(u.nameParts.head).get
         aliases.size match {
           case n if n > 1 =>
-            throw QueryCompilationErrors.ambiguousLateralColumnAlias(u.name, n)
+            throw QueryCompilationErrors.ambiguousLateralColumnAliasError(u.name, n)
           case n if n == 1 && aliases.head.alias.resolved =>
             // Only resolved alias can be the lateral column alias
             // The lateral alias can be a struct and have nested field, need to construct
@@ -164,7 +164,7 @@ object WrapLateralColumnAliasReference extends Rule[LogicalPlan] {
         val aliases = aliasMap.get(nameParts.head).get
         aliases.size match {
           case n if n > 1 =>
-            throw QueryCompilationErrors.ambiguousLateralColumnAlias(nameParts, n)
+            throw QueryCompilationErrors.ambiguousLateralColumnAliasError(nameParts, n)
           case n if n == 1 && aliases.head.alias.resolved =>
             resolveByLateralAlias(nameParts, aliases.head.alias).getOrElse(o)
           case _ => o
@@ -195,6 +195,30 @@ object WrapLateralColumnAliasReference extends Rule[LogicalPlan] {
               wrapLCARef(e, p, aliasMap)
           }
           p.copy(projectList = newProjectList)
+
+        // Implementation notes:
+        // In Aggregate, introducing and wrapping this resolved leaf expression
+        // LateralColumnAliasReference is especially needed because it needs an accurate condition
+        // to trigger adding a Project above and extracting and pushing down aggregate functions
+        // or grouping expressions. Such operation can only be done once. With this
+        // LateralColumnAliasReference, that condition can simply be when the whole Aggregate is
+        // resolved. Otherwise, it can't tell if all aggregate functions are created and
+        // resolved so that it can start the extraction, because the lateral alias reference is
+        // unresolved and can be the argument to functions, blocking the resolution of functions.
+        case agg @ Aggregate(_, aggExprs, _) if agg.childrenResolved
+          && !SimpleAnalyzer.ResolveReferences.containsStar(aggExprs)
+          && aggExprs.exists(_.containsAnyPattern(UNRESOLVED_ATTRIBUTE, OUTER_REFERENCE)) =>
+
+          var aliasMap = CaseInsensitiveMap(Map[String, Seq[AliasEntry]]())
+          val newAggExprs = aggExprs.zipWithIndex.map {
+            case (a: Alias, idx) =>
+              val lcaWrapped = wrapLCARef(a, agg, aliasMap).asInstanceOf[Alias]
+              aliasMap = insertIntoAliasMap(lcaWrapped, idx, aliasMap)
+              lcaWrapped
+            case (e, _) =>
+              wrapLCARef(e, agg, aliasMap)
+          }
+          agg.copy(aggregateExpressions = newAggExprs)
       }
     }
   }