handle aggs too

Author: marmbrus

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

@@ -247,15 +247,13 @@ class Analyzer(catalog: Catalog,
             (oldVersion, newVersion)
 
           // Handle projects that create conflicting aliases.
-          case oldVersion @ Project(projectList, child)
-              if newAliases(projectList).intersect(conflictingAttributes).nonEmpty =>
-            val newVersion =
-              oldVersion.copy(
-                projectList = projectList.map {
-                  case a: Alias => Alias(a.child, a.name)()
-                  case other => other
-                })
-            (oldVersion, newVersion)
+          case oldVersion @ Project(projectList, _)
+              if findAliases(projectList).intersect(conflictingAttributes).nonEmpty =>
+            (oldVersion, oldVersion.copy(projectList = newAliases(projectList)))
+
+          case oldVersion @ Aggregate(_, aggregateExpressions, _)
+              if findAliases(aggregateExpressions).intersect(conflictingAttributes).nonEmpty =>
+            (oldVersion, oldVersion.copy(aggregateExpressions = newAliases(aggregateExpressions)))
         }.head // Only handle first case found, others will be fixed on the next pass.
 
         val attributeRewrites = AttributeMap(oldRelation.output.zip(newRelation.output))
@@ -285,7 +283,14 @@ class Analyzer(catalog: Catalog,
         }
     }
 
-    def newAliases(projectList: Seq[NamedExpression]): AttributeSet = {
+    def newAliases(expressions: Seq[NamedExpression]): Seq[NamedExpression] = {
+      expressions.map {
+        case a: Alias => Alias(a.child, a.name)()
+        case other => other
+      }
+    }
+
+    def findAliases(projectList: Seq[NamedExpression]): AttributeSet = {
       AttributeSet(projectList.collect { case a: Alias => a.toAttribute })
     }
 

sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala

@@ -37,7 +37,7 @@ class SQLQuerySuite extends QueryTest with BeforeAndAfterAll {
   val sqlCtx = TestSQLContext
 
   test("self join with aliases") {
-    val df = Seq(1,2,3).map(i => (i, i.toString)).toDF("int", "str").registerTempTable("df")
+    Seq(1,2,3).map(i => (i, i.toString)).toDF("int", "str").registerTempTable("df")
 
     checkAnswer(
       sql(
@@ -49,6 +49,24 @@ class SQLQuerySuite extends QueryTest with BeforeAndAfterAll {
       Row("1", 1) :: Row("2", 1) :: Row("3", 1) :: Nil)
   }
 
+  test("self join with alias in agg") {
+      Seq(1,2,3)
+        .map(i => (i, i.toString))
+        .toDF("int", "str")
+        .groupBy("str")
+        .agg($"str", count("str").as("strCount"))
+        .registerTempTable("df")
+
+    checkAnswer(
+      sql(
+        """
+          |SELECT x.str, SUM(x.strCount)
+          |FROM df x JOIN df y ON x.str = y.str
+          |GROUP BY x.str
+        """.stripMargin),
+      Row("1", 1) :: Row("2", 1) :: Row("3", 1) :: Nil)
+  }
+
   test("SPARK-4625 support SORT BY in SimpleSQLParser & DSL") {
     checkAnswer(
       sql("SELECT a FROM testData2 SORT BY a"),