fix ut

apache · Jul 16, 2020 · 0950e9a · 0950e9a
1 parent 7465bfa
commit 0950e9a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeSkewedJoin.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeSkewedJoin.scala
@@ -199,9 +199,7 @@ case class OptimizeSkewedJoin(conf: SQLConf) extends Rule[SparkPlan] {
    *    3 tasks separately.
    */
   def optimizeSkewJoin(plan: SparkPlan): SparkPlan = plan.transformUp {
-    case smj @ SortMergeJoinExec(_, _, joinType, _,
-        s1 @ SortExec(_, _, _, _),
-        s2 @ SortExec(_, _, _, _), _)
+    case smj @ SortMergeJoinExec(_, _, joinType, _, s1: SortExec, s2: SortExec, _)
         if supportedJoinTypes.contains(joinType) =>
       // find the shuffleStage from the plan tree
       val leftOpt = findShuffleStage(s1)

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
@@ -660,9 +660,9 @@ class AdaptiveQueryExecSuite
 
         checkSkewJoin(
           "SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2", true)
-        // Additional shuffle introduced, so disable the "OptimizeSkewedJoin" optimization
+        // After patched SPARK-32201, this query won't introduce additional shuffle anymore.
         checkSkewJoin(
-          "SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 GROUP BY key1", false)
+          "SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 GROUP BY key1", true)
       }
     }
   }