Address all comments

Author: wangyum

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

@@ -19,7 +19,7 @@ package org.apache.spark.sql.catalyst.plans.logical
 
 import org.apache.spark.sql.catalyst.expressions.{Alias, Expression, ExpressionSet, NamedExpression}
 import org.apache.spark.sql.catalyst.planning.ExtractEquiJoinKeys
-import org.apache.spark.sql.catalyst.plans.{Inner, LeftSemiOrAnti}
+import org.apache.spark.sql.catalyst.plans.{Inner, LeftOuter, LeftSemiOrAnti, RightOuter}
 
 /**
  * A visitor pattern for traversing a [[LogicalPlan]] tree and propagate the distinct attributes.
@@ -83,10 +83,21 @@ object DistinctKeyVisitor extends LogicalPlanVisitor[Set[ExpressionSet]] {
     p match {
       case Join(_, _, LeftSemiOrAnti(_), _, _) =>
         p.left.distinctKeys
-      case ExtractEquiJoinKeys(Inner, leftKeys, rightKeys, _, _, _, _, _)
-        if p.left.distinctKeys.exists(_.subsetOf(ExpressionSet(leftKeys))) &&
-          p.right.distinctKeys.exists(_.subsetOf(ExpressionSet(rightKeys))) =>
-        Set(ExpressionSet(leftKeys), ExpressionSet(rightKeys))
+      case ExtractEquiJoinKeys(joinType, leftKeys, rightKeys, _, _, left, right, _)
+          if left.distinctKeys.nonEmpty || right.distinctKeys.nonEmpty =>
+        val rightJoinKeySet = ExpressionSet(rightKeys)
+        val leftJoinKeySet = ExpressionSet(leftKeys)
+        joinType match {
+          case Inner if left.distinctKeys.exists(_.subsetOf(leftJoinKeySet)) &&
+            right.distinctKeys.exists(_.subsetOf(rightJoinKeySet)) =>
+            left.distinctKeys ++ right.distinctKeys
+          case Inner | LeftOuter if right.distinctKeys.exists(_.subsetOf(rightJoinKeySet)) =>
+            p.left.distinctKeys
+          case Inner | RightOuter if left.distinctKeys.exists(_.subsetOf(leftJoinKeySet)) =>
+            p.right.distinctKeys
+          case _ =>
+            default(p)
+        }
       case _ => default(p)
     }
   }

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

@@ -150,6 +150,13 @@ class RemoveRedundantAggregatesSuite extends PlanTest {
     comparePlans(optimized, expected)
   }
 
+  test("Remove redundant aggregate - upper has contains foldable expressions") {
+    val originalQuery = x.groupBy('a, 'b)('a, 'b).groupBy('a)('a, TrueLiteral).analyze
+    val correctAnswer = x.groupBy('a)('a, TrueLiteral).analyze
+    val optimized = Optimize.execute(originalQuery)
+    comparePlans(optimized, correctAnswer)
+  }
+
   test("Keep non-redundant aggregate - upper references agg expression") {
     for (agg <- aggregates('b)) {
       val query = relation

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

@@ -104,19 +104,29 @@ class DistinctKeyVisitorSuite extends PlanTest {
       Distinct(t1).join(Distinct(t2), Inner, Some('a === 'x && 'b === 'y && 'c === 'z)),
       Set(ExpressionSet(Seq(a, b, c)), ExpressionSet(Seq(x, y, z))))
 
-    checkDistinctAttributes(t1.join(t2, LeftSemi, Some('a === 'x)),
-      Set.empty)
     checkDistinctAttributes(
-      Distinct(t1).join(Distinct(t2), Inner, Some('a === 'x && 'b === 'y)),
-      Set.empty)
+      Distinct(t1).join(Distinct(t2), LeftOuter, Some('a === 'x && 'b === 'y && 'c === 'z)),
+      Set(ExpressionSet(Seq(a, b, c))))
+
     checkDistinctAttributes(
-      Distinct(t1).join(Distinct(t2), Inner, Some('a === 'x && 'b === 'y && 'c % 5 === 'z % 5)),
-      Set.empty)
-    Seq(LeftOuter, Cross, RightOuter).foreach { joinType =>
+      Distinct(t1).join(Distinct(t2), RightOuter, Some('a === 'x && 'b === 'y && 'c === 'z)),
+      Set(ExpressionSet(Seq(x, y, z))))
+
+    Seq(Inner, Cross, LeftOuter, RightOuter).foreach { joinType =>
+      checkDistinctAttributes(t1.join(t2, joinType, Some('a === 'x)),
+        Set.empty)
+      checkDistinctAttributes(
+        Distinct(t1).join(Distinct(t2), joinType, Some('a === 'x && 'b === 'y)),
+        Set.empty)
       checkDistinctAttributes(
-        Distinct(t1).join(Distinct(t2), joinType, Some('a === 'x && 'b === 'y && 'c === 'z)),
+        Distinct(t1).join(Distinct(t2), joinType,
+          Some('a === 'x && 'b === 'y && 'c % 5 === 'z % 5)),
         Set.empty)
     }
+
+    checkDistinctAttributes(
+      Distinct(t1).join(Distinct(t2), Cross, Some('a === 'x && 'b === 'y && 'c === 'z)),
+      Set.empty)
   }
 
   test("Project's distinct attributes") {