[SPARK-24495][SQL] EnsureRequirement returns worng plan when reordering equal keys

Author: mgaido91

sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala

@@ -17,6 +17,7 @@
 
 package org.apache.spark.sql.execution.exchange
 
+import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
 
 import org.apache.spark.sql.catalyst.expressions._
@@ -227,9 +228,16 @@ case class EnsureRequirements(conf: SQLConf) extends Rule[SparkPlan] {
       currentOrderOfKeys: Seq[Expression]): (Seq[Expression], Seq[Expression]) = {
     val leftKeysBuffer = ArrayBuffer[Expression]()
     val rightKeysBuffer = ArrayBuffer[Expression]()
+    val alreadyUsedIndexes = mutable.Set[Int]()
+    val keysAndIndexes = currentOrderOfKeys.zipWithIndex
 
     expectedOrderOfKeys.foreach(expression => {
-      val index = currentOrderOfKeys.indexWhere(e => e.semanticEquals(expression))
+      val index = keysAndIndexes.find { case (e, idx) =>
+        // As we may have the same key used many times, we need to filter out its occurrence we
+        // have already used.
+        e.semanticEquals(expression) && !alreadyUsedIndexes.contains(idx)
+      }.map(_._2).get
+      alreadyUsedIndexes += index
       leftKeysBuffer.append(leftKeys(index))
       rightKeysBuffer.append(rightKeys(index))
     })

sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala

@@ -679,6 +679,17 @@ class PlannerSuite extends SharedSQLContext {
     }
     assert(rangeExecInZeroPartition.head.outputPartitioning == UnknownPartitioning(0))
   }
+
+  test("SPARK-24495: EnsureRequirements can return wrong plan when reusing the same key in join") {
+    withSQLConf(("spark.sql.shuffle.partitions", "1"),
+      ("spark.sql.constraintPropagation.enabled", "false"),
+      ("spark.sql.autoBroadcastJoinThreshold", "-1")) {
+      val df1 = spark.range(100)
+      val df2 = spark.range(100).select(($"id" * 2).as("b1"), (- $"id").as("b2"))
+      val res = df1.join(df2, $"id" === $"b1" && $"id" === $"b2")
+      assert(res.collect().sameElements(Array(Row(0, 0, 0))))
+    }
+  }
 }
 
 // Used for unit-testing EnsureRequirements