Address all comments

Author: wangyum

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

@@ -48,14 +48,12 @@ object RemoveRedundantAggregates extends Rule[LogicalPlan] with AliasHelper {
         newAggregate
       }
 
-     case agg @ Aggregate(groupingExps, _, child)
-         if agg.groupOnly && child.deterministic &&
-           child.distinctKeys.exists(_.subsetOf(ExpressionSet(groupingExps))) =>
+    case agg @ Aggregate(groupingExps, _, child)
+        if agg.groupOnly && child.distinctKeys.exists(_.subsetOf(ExpressionSet(groupingExps))) =>
       Project(agg.aggregateExpressions, child)
 
     case agg @ Aggregate(groupingExps, aggregateExps, child)
         if aggregateExps.forall(a => a.isInstanceOf[Alias] && a.children.forall(_.foldable)) &&
-          child.deterministic &&
           child.distinctKeys.exists(_.subsetOf(ExpressionSet(groupingExps))) =>
       Project(agg.aggregateExpressions, child)
   }

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

@@ -38,6 +38,7 @@ object DistinctKeyVisitor extends LogicalPlanVisitor[Set[ExpressionSet]] {
         expressionSet.map { expression =>
           expression transform {
             case expr: Expression =>
+              // TODO: Expand distinctKeys for redundant aliases on the same expression
               aliases
                 .collectFirst { case a: Alias if a.child.semanticEquals(expr) => a.toAttribute }
                 .getOrElse(expr)
@@ -60,7 +61,7 @@ object DistinctKeyVisitor extends LogicalPlanVisitor[Set[ExpressionSet]] {
   override def visitDistinct(p: Distinct): Set[ExpressionSet] = Set(ExpressionSet(p.output))
 
   override def visitExcept(p: Except): Set[ExpressionSet] =
-    if (!p.isAll && p.deterministic) Set(ExpressionSet(p.output)) else default(p)
+    if (!p.isAll) Set(ExpressionSet(p.output)) else default(p)
 
   override def visitExpand(p: Expand): Set[ExpressionSet] = default(p)
 
@@ -76,7 +77,7 @@ object DistinctKeyVisitor extends LogicalPlanVisitor[Set[ExpressionSet]] {
   }
 
   override def visitIntersect(p: Intersect): Set[ExpressionSet] = {
-    if (!p.isAll && p.deterministic) Set(ExpressionSet(p.output)) else default(p)
+    if (!p.isAll) Set(ExpressionSet(p.output)) else default(p)
   }
 
   override def visitJoin(p: Join): Set[ExpressionSet] = {
@@ -108,7 +109,7 @@ object DistinctKeyVisitor extends LogicalPlanVisitor[Set[ExpressionSet]] {
 
   override def visitProject(p: Project): Set[ExpressionSet] = {
     if (p.child.distinctKeys.nonEmpty) {
-      projectDistinctKeys(p.child.distinctKeys.map(ExpressionSet(_)), p.projectList)
+      projectDistinctKeys(p.child.distinctKeys, p.projectList)
     } else {
       default(p)
     }

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

@@ -168,13 +168,17 @@ class RemoveRedundantAggregatesSuite extends PlanTest {
     }
   }
 
-  test("Keep non-redundant aggregate - upper references non-deterministic non-grouping") {
+  test("Remove non-redundant aggregate - upper references non-deterministic non-grouping") {
     val query = relation
       .groupBy('a)('a, ('a + rand(0)) as 'c)
       .groupBy('a, 'c)('a, 'c)
       .analyze
+    val expected = relation
+      .groupBy('a)('a, ('a + rand(0)) as 'c)
+      .select('a, 'c)
+      .analyze
     val optimized = Optimize.execute(query)
-    comparePlans(optimized, query)
+    comparePlans(optimized, expected)
   }
 
   test("SPARK-36194: Remove aggregation from left semi/anti join if aggregation the same") {
@@ -283,13 +287,4 @@ class RemoveRedundantAggregatesSuite extends PlanTest {
     val originalQuery = Distinct(x.groupBy('a, 'b)('a, TrueLiteral)).analyze
     comparePlans(Optimize.execute(originalQuery), originalQuery)
   }
-
-  test("SPARK-36194: Negative case: Remove aggregation from contains non-deterministic") {
-    val query = relation
-      .groupBy('a)('a, (count('b) + rand(0)).as("cnt"))
-      .groupBy('a, 'cnt)('a, 'cnt)
-      .analyze
-    val optimized = Optimize.execute(query)
-    comparePlans(optimized, query)
-  }
 }