[SPARK-36183][SQL][FOLLOWUP] Fix push down limit 1 through Aggregate (#861)

* [SPARK-36183][SQL][FOLLOWUP] Fix push down limit 1 through Aggregate

### What changes were proposed in this pull request?

Use `Aggregate.aggregateExpressions` instead of `Aggregate.output` when  pushing down limit 1 through Aggregate.

For example:

```scala
spark.range(10).selectExpr("id % 5 AS a", "id % 5 AS b").write.saveAsTable("t1")
spark.sql("SELECT a, b, a AS alias FROM t1 GROUP BY a, b LIMIT 1").explain(true)
```
Before this pr:
```
== Optimized Logical Plan ==
GlobalLimit 1
+- LocalLimit 1
   +- !Project [a#227L, b#228L, alias#226L]
      +- LocalLimit 1
         +- Relation default.t1[a#227L,b#228L] parquet
```
After this pr:
```
== Optimized Logical Plan ==
GlobalLimit 1
+- LocalLimit 1
   +- Project [a#227L, b#228L, a#227L AS alias#226L]
      +- LocalLimit 1
         +- Relation default.t1[a#227L,b#228L] parquet
```

### Why are the changes needed?

Fix bug.

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?

Unit test.

Closes #35286 from wangyum/SPARK-36183-2.

Authored-by: Yuming Wang <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>

(cherry picked from commit 9b12571)

Author: wangyum

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

@@ -548,9 +548,9 @@ object LimitPushDown extends Rule[LogicalPlan] {
 
     // Push down limit 1 through Aggregate and turn Aggregate into Project if it is group only.
     case Limit(le @ IntegerLiteral(1), a: Aggregate) if a.groupOnly =>
-      Limit(le, Project(a.output, LocalLimit(le, a.child)))
+      Limit(le, Project(a.aggregateExpressions, LocalLimit(le, a.child)))
     case Limit(le @ IntegerLiteral(1), p @ Project(_, a: Aggregate)) if a.groupOnly =>
-      Limit(le, p.copy(child = Project(a.output, LocalLimit(le, a.child))))
+      Limit(le, p.copy(child = Project(a.aggregateExpressions, LocalLimit(le, a.child))))
   }
 }
 

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/LimitPushdownSuite.scala

@@ -182,6 +182,19 @@ class LimitPushdownSuite extends PlanTest {
       Optimize.execute(x.groupBy("x.a".attr)("x.a".attr).select("x.a".attr).limit(1).analyze),
       LocalLimit(1, x).select("x.a".attr).select("x.a".attr).limit(1).analyze)
 
+    comparePlans(
+      Optimize.execute(x.union(y).groupBy("x.a".attr)("x.a".attr).limit(1).analyze),
+      LocalLimit(1, LocalLimit(1, x).union(
+        LocalLimit(1, y.select('a.as("a"), 'b.as("b"), 'c.as("c"))))).select("x.a".attr)
+        .limit(1).analyze)
+
+    comparePlans(
+      Optimize.execute(
+        x.groupBy("x.a".attr)("x.a".attr)
+          .select("x.a".attr.as("a1"), "x.a".attr.as("a2")).limit(1).analyze),
+      LocalLimit(1, x).select("x.a".attr)
+        .select("x.a".attr.as("a1"), "x.a".attr.as("a2")).limit(1).analyze)
+
     // No push down
     comparePlans(
       Optimize.execute(x.groupBy("x.a".attr)("x.a".attr).limit(2).analyze),