add ut

apache · Jul 20, 2020 · 8aee8a5 · 8aee8a5
1 parent bc454b6
commit 8aee8a5
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/...catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/EliminateSortsSuite.scala b/...catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/EliminateSortsSuite.scala
@@ -24,6 +24,7 @@ import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.dsl.plans._
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate.{MaxBy, MinBy}
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules._
@@ -319,4 +320,15 @@ class EliminateSortsSuite extends PlanTest {
     val correctAnswer = PushDownOptimizer.execute(noOrderByPlan.analyze)
     comparePlans(optimized, correctAnswer)
   }
+
+  test("SPARK-32360: Add MaxMinBy to support eliminate sorts") {
+    Seq(MaxBy(Symbol("a"), Symbol("b")), MinBy(Symbol("a"), Symbol("b"))).foreach { agg =>
+      val projectPlan = testRelation.select(Symbol("a"), Symbol("b"))
+      val unnecessaryOrderByPlan = projectPlan.orderBy(Symbol("a").asc)
+      val groupByPlan = unnecessaryOrderByPlan.groupBy(Symbol("a"))(agg)
+      val optimized = Optimize.execute(groupByPlan.analyze)
+      val correctAnswer = projectPlan.groupBy(Symbol("a"))(agg).analyze
+      comparePlans(optimized, correctAnswer)
+    }
+  }
 }