Allow duplicate angostic aggs

Author: tanelk

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

@@ -358,7 +358,7 @@ object EliminateDistinct extends Rule[LogicalPlan] {
       ae.copy(isDistinct = false)
   }
 
-  private def isDuplicateAgnostic(af: AggregateFunction): Boolean = af match {
+  def isDuplicateAgnostic(af: AggregateFunction): Boolean = af match {
     case _: Max => true
     case _: Min => true
     case _: BitAndAgg => true

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

@@ -48,7 +48,9 @@ object RemoveRedundantAggregates extends Rule[LogicalPlan] with AliasHelper {
   }
 
   private def isLowerRedundant(upper: Aggregate, lower: Aggregate): Boolean = {
-    val upperHasNoAggregateExpressions = !upper.aggregateExpressions.exists(isAggregate)
+    val upperHasNoDuplicateSensitiveAgg = !upper
+      .aggregateExpressions
+      .exists(isDuplicateSensitiveAggregate)
 
     lazy val upperRefsOnlyDeterministicNonAgg = upper.references.subsetOf(AttributeSet(
       lower
@@ -58,11 +60,18 @@ object RemoveRedundantAggregates extends Rule[LogicalPlan] with AliasHelper {
         .map(_.toAttribute)
     ))
 
-    upperHasNoAggregateExpressions && upperRefsOnlyDeterministicNonAgg
+    upperHasNoDuplicateSensitiveAgg && upperRefsOnlyDeterministicNonAgg
   }
 
   private def isAggregate(expr: Expression): Boolean = {
     expr.find(e => e.isInstanceOf[AggregateExpression] ||
       PythonUDF.isGroupedAggPandasUDF(e)).isDefined
   }
+
+  private def isDuplicateSensitiveAggregate(expr: Expression): Boolean = {
+    expr.find {
+      case ae: AggregateExpression => !EliminateDistinct.isDuplicateAgnostic(ae.aggregateFunction)
+      case e => PythonUDF.isGroupedAggPandasUDF(e)
+    }.isDefined
+  }
 }

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

@@ -126,7 +126,7 @@ class RemoveRedundantAggregatesSuite extends PlanTest {
     comparePlans(optimized, expected)
   }
 
-  test("Keep non-redundant aggregate - upper has agg expression") {
+  test("Keep non-redundant aggregate - upper has duplicate sensitive agg expression") {
     val relation = LocalRelation('a.int, 'b.int)
     for (agg <- aggregates('b)) {
       val query = relation
@@ -139,6 +139,20 @@ class RemoveRedundantAggregatesSuite extends PlanTest {
     }
   }
 
+  test("Remove redundant aggregate - upper has duplicate agnostic agg expression") {
+    val relation = LocalRelation('a.int, 'b.int)
+    val query = relation
+      .groupBy('a, 'b)('a, 'b)
+      // The max does not change if there are duplicate values
+      .groupBy('a)('a, max('b))
+      .analyze
+    val expected = relation
+      .groupBy('a)('a, max('b))
+      .analyze
+    val optimized = Optimize.execute(query)
+    comparePlans(optimized, expected)
+  }
+
   test("Keep non-redundant aggregate - upper references agg expression") {
     val relation = LocalRelation('a.int, 'b.int)
     for (agg <- aggregates('b)) {