[SPARK-46378][SQL] Still remove Sort after converting Aggregate to Project

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

This is a follow-up of apache#33397 to avoid sub-optimal plans. After converting `Aggregate` to `Project`, there is information lost: `Aggregate` doesn't care about the data order of inputs, but `Project` cares. `EliminateSorts` can remove `Sort` below `Aggregate`, but it doesn't work anymore if we convert `Aggregate` to `Project`.

This PR fixes this issue by tagging the `Project` to be order-irrelevant if it's converted from `Aggregate`. Then `EliminateSorts` optimizes the tagged `Project`.

### Why are the changes needed?

avoid sub-optimal plans

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

No
### How was this patch tested?

new test

### Was this patch authored or co-authored using generative AI tooling?

No

Closes apache#44310 from cloud-fan/sort.

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

Author: cloud-fan

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/dsl/package.scala

@@ -404,6 +404,8 @@ package object dsl {
 
       def limit(limitExpr: Expression): LogicalPlan = Limit(limitExpr, logicalPlan)
 
+      def localLimit(limitExpr: Expression): LogicalPlan = LocalLimit(limitExpr, logicalPlan)
+
       def offset(offsetExpr: Expression): LogicalPlan = Offset(offsetExpr, logicalPlan)
 
       def join(

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

@@ -757,7 +757,9 @@ object LimitPushDown extends Rule[LogicalPlan] {
       LocalLimit(exp, project.copy(child = pushLocalLimitThroughJoin(exp, join)))
     // 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.aggregateExpressions, LocalLimit(le, a.child)))
+      val project = Project(a.aggregateExpressions, LocalLimit(le, a.child))
+      project.setTagValue(Project.dataOrderIrrelevantTag, ())
+      Limit(le, project)
     case Limit(le @ IntegerLiteral(1), p @ Project(_, a: Aggregate)) if a.groupOnly =>
       Limit(le, p.copy(child = Project(a.aggregateExpressions, LocalLimit(le, a.child))))
     // Merge offset value and limit value into LocalLimit and pushes down LocalLimit through Offset.
@@ -1563,6 +1565,8 @@ object EliminateSorts extends Rule[LogicalPlan] {
         right = recursiveRemoveSort(originRight, true))
     case g @ Aggregate(_, aggs, originChild) if isOrderIrrelevantAggs(aggs) =>
       g.copy(child = recursiveRemoveSort(originChild, true))
+    case p: Project if p.getTagValue(Project.dataOrderIrrelevantTag).isDefined =>
+      p.copy(child = recursiveRemoveSort(p.child, true))
   }
 
   /**

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala

@@ -98,6 +98,9 @@ case class Project(projectList: Seq[NamedExpression], child: LogicalPlan)
 
 object Project {
   val hiddenOutputTag: TreeNodeTag[Seq[Attribute]] = TreeNodeTag[Seq[Attribute]]("hidden_output")
+  // Project with this tag means it doesn't care about the data order of its input. We only set
+  // this tag when the Project was converted from grouping-only Aggregate.
+  val dataOrderIrrelevantTag: TreeNodeTag[Unit] = TreeNodeTag[Unit]("data_order_irrelevant")
 
   def matchSchema(plan: LogicalPlan, schema: StructType, conf: SQLConf): Project = {
     assert(plan.resolved)

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

@@ -478,4 +478,16 @@ class EliminateSortsSuite extends AnalysisTest {
 
     comparePlans(Optimize.execute(originalPlan.analyze), correctAnswer.analyze)
   }
+
+  test("SPARK-46378: Still remove Sort after converting Aggregate to Project") {
+    val originalPlan = testRelation.orderBy($"a".asc)
+      .groupBy($"a")($"a")
+      .limit(1)
+
+    val correctAnswer = testRelation.localLimit(1)
+      .select($"a")
+      .limit(1)
+
+    comparePlans(Optimize.execute(originalPlan.analyze), correctAnswer.analyze)
+  }
 }