Add a Project node in case an aggregation is relying on aliases missi…

…ng in the child plan
mbautin · Jan 29, 2016 · 2e99064 · 2e99064
1 parent 5267564
commit 2e99064
Showing 1 changed file with 19 additions and 3 deletions.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
@@ -195,6 +195,22 @@ private[sql] abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
           }.asInstanceOf[NamedExpression]
         }
 
+        var exprIdsInChild = Set[ExprId]()
+        child.transformAllExpressions {
+          case expr: NamedExpression =>
+            exprIdsInChild += expr.exprId
+            expr
+          case other => other
+        }
+
+        val namedGroupingExpressionsUnknownToChild =
+          namedGroupingExpressions.map(_._2).filterNot { ne => exprIdsInChild.contains(ne.exprId) }
+        val childWithExtraAliases = if (namedGroupingExpressionsUnknownToChild.isEmpty) {
+          child
+        } else {
+          logical.Project(child.output ++ namedGroupingExpressionsUnknownToChild, child)
+        }
+
         val aggregateOperator =
           if (aggregateExpressions.map(_.aggregateFunction).exists(!_.supportsPartial)) {
             if (functionsWithDistinct.nonEmpty) {
@@ -206,23 +222,23 @@ private[sql] abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
                 aggregateExpressions,
                 aggregateFunctionToAttribute,
                 rewrittenResultExpressions,
-                planLater(child))
+                planLater(childWithExtraAliases))
             }
           } else if (functionsWithDistinct.isEmpty) {
             aggregate.Utils.planAggregateWithoutDistinct(
               namedGroupingExpressions.map(_._2),
               aggregateExpressions,
               aggregateFunctionToAttribute,
               rewrittenResultExpressions,
-              planLater(child))
+              planLater(childWithExtraAliases))
           } else {
             aggregate.Utils.planAggregateWithOneDistinct(
               namedGroupingExpressions.map(_._2),
               functionsWithDistinct,
               functionsWithoutDistinct,
               aggregateFunctionToAttribute,
               rewrittenResultExpressions,
-              planLater(child))
+              planLater(childWithExtraAliases))
           }
 
         aggregateOperator