handle Project+Aggregate with `mergeExpressions(), add more comment.

Author: peter-toth

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

@@ -1206,7 +1206,7 @@ object CollapseProject extends Rule[LogicalPlan] with AliasHelper {
       newPythonUDFEvalTypesInUpperProjects,
       pythonUDFArrowFallbackOnUDT)) match {
       case p1 @ Project(_, p2: Project) =>
-        mergeProjectExpressions(
+        mergeExpressions(
           p1.projectList,
           p2.projectList,
           alwaysInline,
@@ -1216,16 +1216,17 @@ object CollapseProject extends Rule[LogicalPlan] with AliasHelper {
             case (newUpper, newLower) =>
               p1.copy(projectList = newUpper, child = p2.copy(projectList = newLower))
         }
-      case p @ Project(_, agg: Aggregate)
-          if canCollapseExpressions(
-            p.projectList,
-            getAliasMap(agg.aggregateExpressions),
-            alwaysInline,
-            newPythonUDFEvalTypesInUpperProjects,
-            pythonUDFArrowFallbackOnUDT)
-          && canCollapseAggregate(p, agg) =>
-        agg.copy(aggregateExpressions = buildCleanedProjectList(
-          p.projectList, agg.aggregateExpressions))
+      case p @ Project(_, agg: Aggregate) if canCollapseAggregate(p, agg) =>
+        mergeExpressions(
+          p.projectList,
+          agg.aggregateExpressions,
+          alwaysInline,
+          newPythonUDFEvalTypesInUpperProjects,
+          pythonUDFArrowFallbackOnUDT) match {
+            case (Seq(), merged) => agg.copy(aggregateExpressions = merged)
+            case (newUpper, newLower) =>
+              p.copy(projectList = newUpper, child = agg.copy(aggregateExpressions = newLower))
+        }
       case Project(l1, g @ GlobalLimit(_, limit @ LocalLimit(_, p2 @ Project(l2, _))))
         if isRenaming(l1, l2) =>
         val newProjectList = buildCleanedProjectList(l1, l2)
@@ -1275,7 +1276,7 @@ object CollapseProject extends Rule[LogicalPlan] with AliasHelper {
     case other => isCheap(other)
   }
 
-  private def mergeProjectExpressions(
+  private def mergeExpressions(
       consumers: Seq[NamedExpression],
       producers: Seq[NamedExpression],
       alwaysInline: Boolean,
@@ -1291,9 +1292,19 @@ object CollapseProject extends Rule[LogicalPlan] with AliasHelper {
       .groupMap(_._1)(_._2)
       .view.mapValues(v => v.size -> ExpressionSet(v)))
 
+    // Split the producers from the lower node to 4 categories:
+    // - `neverInlines` contains producer expressions that shouldn't be inlined.
+    //    These include non-deterministic expressions or expensive ones that are referenced multiple
+    //    times.
+    // - `mustInlines` contains expressions with Python UDFs that must be inlined into the upper
+    //    node to avoid performance issues, or expressions with aggregate nodes.
+    // - `maybeInlines` contains expressions that might make sense to inline, such as expressions
+    //    that are used only once, or are cheap to inline.
+    //    But we need to take into account the side effect of adding new pass-through attributes to
+    //    the lover node, which can make the node much wider than it was originally.
+    val neverInlines = ListBuffer.empty[NamedExpression]
     val mustInlines = ListBuffer.empty[NamedExpression]
     val maybeInlines = ListBuffer.empty[NamedExpression]
-    val neverInlines = ListBuffer.empty[NamedExpression]
     val others = ListBuffer.empty[NamedExpression]
     producers.foreach {
       case a: Alias =>
@@ -1306,7 +1317,7 @@ object CollapseProject extends Rule[LogicalPlan] with AliasHelper {
             case _ => false
           }
 
-          if (!a.deterministic) {
+          if (!a.child.deterministic || AggregateExpression.containsAggregate(a.child)) {
             neverInlines += a
           } else if (alwaysInline || containsUDF) {
             mustInlines += a
@@ -1321,11 +1332,15 @@ object CollapseProject extends Rule[LogicalPlan] with AliasHelper {
     }
 
     if (neverInlines.isEmpty) {
+      // If `neverInlines` is empty then we can collapse the nodes into one.
       (Seq.empty, buildCleanedProjectList(consumers, producers))
     } else if (mustInlines.isEmpty) {
-      // Let's keep `maybeInlines` in the lower node for now.
+      // Otherwise we can't collapse the nodes into one, but if `mustInlines` is empty then we can
+      // keep `maybeInlines` in the lower node for now, so there is no change to the nodes.
       (consumers, producers)
     } else {
+      // If both `neverInlines` and `mustInlines` are not empty, then inline `mustInlines` and add
+      // new pass-through attributes to the lower node.
       val newConsumers = buildCleanedProjectList(consumers, mustInlines)
       val passthroughAttributes = AttributeSet(others.flatMap(_.references))
       val newPassthroughAttributes =