move logic to the beginning of optimization, simplify test

Author: peter-toth

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

@@ -147,6 +147,7 @@ abstract class Optimizer(catalogManager: CatalogManager)
       EliminateView,
       ReplaceExpressions,
       RewriteNonCorrelatedExists,
+      WrapGroupingExpressions,
       ComputeCurrentTime,
       GetCurrentDatabaseAndCatalog(catalogManager)) ::
     //////////////////////////////////////////////////////////////////////////////////////////
@@ -870,19 +871,8 @@ object CollapseProject extends Rule[LogicalPlan] with AliasHelper {
       if (haveCommonNonDeterministicOutput(p.projectList, agg.aggregateExpressions)) {
         p
       } else {
-        val complexGroupingExpressions =
-          ExpressionSet(agg.groupingExpressions.filter(_.children.nonEmpty))
-
-        def wrapGroupingExpression(e: Expression): Expression = e match {
-          case _: AggregateExpression => e
-          case _ if complexGroupingExpressions.contains(e) => GroupingExpression(e)
-          case _ => e.mapChildren(wrapGroupingExpression)
-        }
-
-        val wrappedAggregateExpressions =
-          agg.aggregateExpressions.map(wrapGroupingExpression(_).asInstanceOf[NamedExpression])
-        agg.copy(aggregateExpressions =
-          buildCleanedProjectList(p.projectList, wrappedAggregateExpressions))
+        agg.copy(aggregateExpressions = buildCleanedProjectList(
+          p.projectList, agg.aggregateExpressions))
       }
     case Project(l1, g @ GlobalLimit(_, limit @ LocalLimit(_, p2 @ Project(l2, _))))
         if isRenaming(l1, l2) =>

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

@@ -66,6 +66,34 @@ object RewriteNonCorrelatedExists extends Rule[LogicalPlan] {
   }
 }
 
+/**
+ * Wrap complex grouping expression in aggregate expressions without aggregate function into
+ * `GroupingExpression` nodes so as to avoid further optimizations between the expression and its
+ * parent.
+ *
+ * This is required as further optimizations could change the grouping expression and so make the
+ * aggregate expression invalid.
+ */
+object WrapGroupingExpressions extends Rule[LogicalPlan] {
+  override def apply(plan: LogicalPlan): LogicalPlan = {
+    plan transform {
+      case a: Aggregate =>
+        val complexGroupingExpressions =
+          ExpressionSet(a.groupingExpressions.filter(_.children.nonEmpty))
+
+        def wrapGroupingExpression(e: Expression): Expression = e match {
+          case _: GroupingExpression => e
+          case _: AggregateExpression => e
+          case _ if complexGroupingExpressions.contains(e) => GroupingExpression(e)
+          case _ => e.mapChildren(wrapGroupingExpression)
+        }
+
+        a.copy(aggregateExpressions =
+          a.aggregateExpressions.map(wrapGroupingExpression(_).asInstanceOf[NamedExpression]))
+    }
+  }
+}
+
 /**
  * Computes the current date and time to make sure we return the same result in a single query.
  */

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

@@ -4123,12 +4123,9 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark
 
       val df = spark.sql(
         """
-          |SELECT not(id), c
-          |FROM (
-          |  SELECT t.id IS NULL AS id, count(*) AS c
-          |  FROM t
-          |  GROUP BY t.id IS NULL
-          |) t
+          |SELECT not(t.id IS NULL), count(*) AS c
+          |FROM t
+          |GROUP BY t.id IS NULL
           |""".stripMargin)
       checkAnswer(df, Row(true, 3) :: Row(false, 2) :: Nil)
     }