[SPARK-36194][SQL][FOLLOWUP] Propagate distinct keys more precisely

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

This PR is a followup of #35779 , to propagate distinct keys more precisely in 2 cases:
1. For `LIMIT 1`, each output attribute is a distinct key, not the entire tuple.
2. For aggregate, we can still propagate distinct keys from child.

### Why are the changes needed?

make the optimization cover more cases

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

no

### How was this patch tested?

new tests

Closes #36100 from cloud-fan/followup.

Authored-by: Wenchen Fan <wenchen@databricks.com>
Signed-off-by: Yuming Wang <yumwang@ebay.com>

Author: cloud-fan

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

@@ -50,11 +50,27 @@ object DistinctKeyVisitor extends LogicalPlanVisitor[Set[ExpressionSet]] {
     }.filter(_.nonEmpty)
   }
 
+  /**
+   * Add a new ExpressionSet S into distinctKeys D.
+   * To minimize the size of D:
+   * 1. If there is a subset of S in D, return D.
+   * 2. Otherwise, remove all the ExpressionSet containing S from D, and add the new one.
+   */
+  private def addDistinctKey(
+      keys: Set[ExpressionSet],
+      newExpressionSet: ExpressionSet): Set[ExpressionSet] = {
+    if (keys.exists(_.subsetOf(newExpressionSet))) {
+      keys
+    } else {
+      keys.filterNot(s => newExpressionSet.subsetOf(s)) + newExpressionSet
+    }
+  }
+
   override def default(p: LogicalPlan): Set[ExpressionSet] = Set.empty[ExpressionSet]
 
   override def visitAggregate(p: Aggregate): Set[ExpressionSet] = {
     val groupingExps = ExpressionSet(p.groupingExpressions) // handle group by a, a
-    projectDistinctKeys(Set(groupingExps), p.aggregateExpressions)
+    projectDistinctKeys(addDistinctKey(p.child.distinctKeys, groupingExps), p.aggregateExpressions)
   }
 
   override def visitDistinct(p: Distinct): Set[ExpressionSet] = Set(ExpressionSet(p.output))
@@ -70,7 +86,7 @@ object DistinctKeyVisitor extends LogicalPlanVisitor[Set[ExpressionSet]] {
 
   override def visitGlobalLimit(p: GlobalLimit): Set[ExpressionSet] = {
     p.maxRows match {
-      case Some(value) if value <= 1 => Set(ExpressionSet(p.output))
+      case Some(value) if value <= 1 => p.output.map(attr => ExpressionSet(Seq(attr))).toSet
       case _ => p.child.distinctKeys
     }
   }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/logical/DistinctKeyVisitorSuite.scala

@@ -66,6 +66,10 @@ class DistinctKeyVisitorSuite extends PlanTest {
       Set(ExpressionSet(Seq(a)), ExpressionSet(Seq(d.toAttribute))))
     checkDistinctAttributes(t1.groupBy(f.child, $"b")(f, $"b", sum($"c")),
       Set(ExpressionSet(Seq(f.toAttribute, b))))
+
+    // Aggregate should also propagate distinct keys from child
+    checkDistinctAttributes(t1.limit(1).groupBy($"a", $"b")($"a", $"b"),
+      Set(ExpressionSet(Seq(a)), ExpressionSet(Seq(b))))
   }
 
   test("Distinct's distinct attributes") {
@@ -86,7 +90,8 @@ class DistinctKeyVisitorSuite extends PlanTest {
   test("Limit's distinct attributes") {
     checkDistinctAttributes(Distinct(t1).limit(10), Set(ExpressionSet(Seq(a, b, c))))
     checkDistinctAttributes(LocalLimit(10, Distinct(t1)), Set(ExpressionSet(Seq(a, b, c))))
-    checkDistinctAttributes(t1.limit(1), Set(ExpressionSet(Seq(a, b, c))))
+    checkDistinctAttributes(t1.limit(1),
+      Set(ExpressionSet(Seq(a)), ExpressionSet(Seq(b)), ExpressionSet(Seq(c))))
   }
 
   test("Intersect's distinct attributes") {