[SPARK-18137][SQL]Fix RewriteDistinctAggregates UnresolvedException when a UDAF has a foldable TypeCheck

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregates.scala

@@ -23,7 +23,7 @@ import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Expand, LogicalPl
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.types.IntegerType
 
-/**
+/*
  * This rule rewrites an aggregate query with distinct aggregations into an expanded double
  * aggregation in which the regular aggregation expressions and every distinct clause is aggregated
  * in a separate group. The results are then combined in a second aggregate.
@@ -115,9 +115,19 @@ object RewriteDistinctAggregates extends Rule[LogicalPlan] {
     }
 
     // Extract distinct aggregate expressions.
-    val distinctAggGroups = aggExpressions
-      .filter(_.isDistinct)
-      .groupBy(_.aggregateFunction.children.toSet)
+    val distinctAggGroups = aggExpressions.filter(_.isDistinct).groupBy{
+      e =>
+        if (e.aggregateFunction.children.exists(!_.foldable)) {
+          // Only expand the unfoldable children
+          e.aggregateFunction.children.filter(!_.foldable).toSet
+        } else {
+          // If aggregateFunction's children are all foldable
+          // we must expand at least one of the children (here we take the first child),
+          // or If we don't, we will get the wrong result, for example:
+          // count(distinct 1) will be explained to count(1) after the rewrite function.
+          e.aggregateFunction.children.take(1).toSet
+        }
+    }
 
     // Check if the aggregates contains functions that do not support partial aggregation.
     val existsNonPartial = aggExpressions.exists(!_.aggregateFunction.supportsPartial)
@@ -134,10 +144,10 @@ object RewriteDistinctAggregates extends Rule[LogicalPlan] {
 
       // Functions used to modify aggregate functions and their inputs.
       def evalWithinGroup(id: Literal, e: Expression) = If(EqualTo(gid, id), e, nullify(e))
-      def patchAggregateFunctionChildren(
-          af: AggregateFunction)(
-          attrs: Expression => Expression): AggregateFunction = {
-        af.withNewChildren(af.children.map(attrs)).asInstanceOf[AggregateFunction]
+      def patchAggregateFunctionChildren(af: AggregateFunction)(
+        attrs: Expression => Option[Expression]): AggregateFunction = {
+        val newChildren = af.children.map(c => attrs(c).getOrElse(c))
+        af.withNewChildren(newChildren).asInstanceOf[AggregateFunction]
       }
 
       // Setup unique distinct aggregate children.
@@ -161,7 +171,7 @@ object RewriteDistinctAggregates extends Rule[LogicalPlan] {
           val operators = expressions.map { e =>
             val af = e.aggregateFunction
             val naf = patchAggregateFunctionChildren(af) { x =>
-              evalWithinGroup(id, distinctAggChildAttrLookup(x))
+              distinctAggChildAttrLookup.get(x).map(evalWithinGroup(id, _))
             }
             (e, e.copy(aggregateFunction = naf, isDistinct = false))
           }
@@ -170,16 +180,20 @@ object RewriteDistinctAggregates extends Rule[LogicalPlan] {
       }
 
       // Setup expand for the 'regular' aggregate expressions.
-      val regularAggExprs = aggExpressions.filter(!_.isDistinct)
-      val regularAggChildren = regularAggExprs.flatMap(_.aggregateFunction.children).distinct
+      // only expand unfoldable children
+      val regularAggExprs = aggExpressions
+        .filter(e => !e.isDistinct && e.children.exists(!_.foldable))
+      val regularAggChildren = regularAggExprs
+        .flatMap(_.aggregateFunction.children.filter(!_.foldable))
+        .distinct
       val regularAggChildAttrMap = regularAggChildren.map(expressionAttributePair)
 
       // Setup aggregates for 'regular' aggregate expressions.
       val regularGroupId = Literal(0)
       val regularAggChildAttrLookup = regularAggChildAttrMap.toMap
       val regularAggOperatorMap = regularAggExprs.map { e =>
         // Perform the actual aggregation in the initial aggregate.
-        val af = patchAggregateFunctionChildren(e.aggregateFunction)(regularAggChildAttrLookup)
+        val af = patchAggregateFunctionChildren(e.aggregateFunction)(regularAggChildAttrLookup.get)
         val operator = Alias(e.copy(aggregateFunction = af), e.sql)()
 
         // Select the result of the first aggregate in the last aggregate.
@@ -237,8 +251,8 @@ object RewriteDistinctAggregates extends Rule[LogicalPlan] {
 
       // Construct the second aggregate
       val transformations: Map[Expression, Expression] =
-        (distinctAggOperatorMap.flatMap(_._2) ++
-          regularAggOperatorMap.map(e => (e._1, e._3))).toMap
+      (distinctAggOperatorMap.flatMap(_._2) ++
+        regularAggOperatorMap.map(e => (e._1, e._3))).toMap
 
       val patchedAggExpressions = a.aggregateExpressions.map { e =>
         e.transformDown {
@@ -261,9 +275,9 @@ object RewriteDistinctAggregates extends Rule[LogicalPlan] {
   private def nullify(e: Expression) = Literal.create(null, e.dataType)
 
   private def expressionAttributePair(e: Expression) =
-    // We are creating a new reference here instead of reusing the attribute in case of a
-    // NamedExpression. This is done to prevent collisions between distinct and regular aggregate
-    // children, in this case attribute reuse causes the input of the regular aggregate to bound to
-    // the (nulled out) input of the distinct aggregate.
+  // We are creating a new reference here instead of reusing the attribute in case of a
+  // NamedExpression. This is done to prevent collisions between distinct and regular aggregate
+  // children, in this case attribute reuse causes the input of the regular aggregate to bound to
+  // the (nulled out) input of the distinct aggregate.
     e -> AttributeReference(e.sql, e.dataType, nullable = true)()
 }

sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveUDFSuite.scala

@@ -150,6 +150,24 @@ class HiveUDFSuite extends QueryTest with TestHiveSingleton with SQLTestUtils {
   }
 
   test("Generic UDAF aggregates") {
+    checkAnswer(sql("SELECT percentile_approx(2, 0.99999), " +
+      "sum(distinct 1), count(distinct 1,2,3,4) FROM src LIMIT 1"),
+      sql("SELECT 2, 1, 1 FROM src LIMIT 1")
+        .collect().toSeq)
+
+    checkAnswer(sql("SELECT ceiling(percentile_approx(distinct key, 0.99999))" +
+      ", count(distinct key), sum(distinct key), " +
+      "count(distinct 1), sum(distinct 1), sum(1) FROM src LIMIT 1"),
+            sql("SELECT max(key), count(distinct key), sum(distinct key)," +
+              " 1, 1, sum(1) FROM src LIMIT 1")
+              .collect().toSeq)
+
+    checkAnswer(sql("SELECT ceiling(percentile_approx(distinct key, 0.9 + 0.09999))" +
+      ", count(distinct key), sum(distinct key), " +
+      "count(distinct 1), sum(distinct 1), sum(1) FROM src LIMIT 1"),
+      sql("SELECT max(key), count(distinct key), sum(distinct key), 1, 1, sum(1) FROM src LIMIT 1")
+        .collect().toSeq)
+
     checkAnswer(sql("SELECT ceiling(percentile_approx(key, 0.99999D)) FROM src LIMIT 1"),
       sql("SELECT max(key) FROM src LIMIT 1").collect().toSeq)