Trims aliases when resolving and checking aggregate expressions

apache · Jan 12, 2015 · 443538d · 443538d
1 parent aff49a3
commit 443538d
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 8 deletions.
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
@@ -209,10 +209,10 @@ class Analyzer(catalog: Catalog,
 
           aggregateExprs.find { e =>
             !isValidAggregateExpression(e.transform {
-              // Should trim aliases around `GetField`s. These aliases are introduced while
-              // resolving struct field accesses, because `GetField` is not a `NamedExpression`.
-              // (Should we just turn `GetField` into a `NamedExpression`?)
-              case Alias(g: GetField, _) => g
+              // Should trim aliases. These aliases can only be introduced while resolving unnamed
+              // expressions like `GetField` and UDF calls, because GROUP BY clause doesn't allow
+              // aliasing.
+              case a: Alias => a.child
             })
           }.foreach { e =>
             throw new TreeNodeException(plan, s"Expression not in GROUP BY: $e")

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala
@@ -153,11 +153,11 @@ object PartialAggregation {
             partialEvaluations(new TreeNodeRef(e)).finalEvaluation
 
           case e: Expression =>
-            // Should trim aliases around `GetField`s. These aliases are introduced while
-            // resolving struct field accesses, because `GetField` is not a `NamedExpression`.
-            // (Should we just turn `GetField` into a `NamedExpression`?)
             namedGroupingExpressions
-              .get(e.transform { case Alias(g: GetField, _) => g })
+              // Should trim aliases. These aliases can only be introduced while resolving unnamed
+              // expressions like `GetField` and UDF calls, because GROUP BY clause doesn't allow
+              // aliasing.
+              .get(e.transform { case a: Alias => a.child })
               .map(_.toAttribute)
               .getOrElse(e)
         }).asInstanceOf[Seq[NamedExpression]]

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
@@ -997,6 +997,17 @@ class SQLQuerySuite extends QueryTest with BeforeAndAfterAll {
     dropTempTable("data")
   }
 
+  test("SPARK-4296 Grouping field with UDF as sub expression") {
+    registerFunction("triple", (_: Int) * 3)
+    jsonRDD(sparkContext.makeRDD("""{"a": 1}""" :: Nil)).registerTempTable("data")
+    checkAnswer(sql("SELECT triple(a) FROM data GROUP BY triple(a)"), 3)
+    dropTempTable("data")
+
+    jsonRDD(sparkContext.makeRDD("""{"a": 1}""" :: Nil)).registerTempTable("data")
+    checkAnswer(sql("SELECT triple(a) + 1 FROM data GROUP BY triple(a) + 1"), 4)
+    dropTempTable("data")
+  }
+
   test("SPARK-4432 Fix attribute reference resolution error when using ORDER BY") {
     checkAnswer(
       sql("SELECT a + b FROM testData2 ORDER BY a"),