Address all new comments

Author: c21

sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashedRelation.scala

@@ -97,17 +97,18 @@ private[execution] object HashedRelation {
   /**
    * Create a HashedRelation from an Iterator of InternalRow.
    *
-   * @param isLookupAware reserve one extra boolean in value to track if value being looked up
-   * @param value the expressions for value inserted into HashedRelation
+   * @param canMarkRowLookedUp Reserve one extra boolean in value to track if value being looked up.
+   *                           This is only used for full outer shuffled hash join.
+   * @param valueExprs The expressions for value inserted into HashedRelation.
    */
   def apply(
       input: Iterator[InternalRow],
       key: Seq[Expression],
       sizeEstimate: Int = 64,
       taskMemoryManager: TaskMemoryManager = null,
       isNullAware: Boolean = false,
-      isLookupAware: Boolean = false,
-      value: Option[Seq[Expression]] = None): HashedRelation = {
+      canMarkRowLookedUp: Boolean = false,
+      valueExprs: Option[Seq[Expression]] = None): HashedRelation = {
     val mm = Option(taskMemoryManager).getOrElse {
       new TaskMemoryManager(
         new UnifiedMemoryManager(
@@ -120,12 +121,13 @@ private[execution] object HashedRelation {
 
     if (isNullAware && !input.hasNext) {
       EmptyHashedRelation
-    } else if (key.length == 1 && key.head.dataType == LongType && !isLookupAware) {
-      // NOTE: LongHashedRelation cannot support isLookupAware as it cannot
+    } else if (key.length == 1 && key.head.dataType == LongType && !canMarkRowLookedUp) {
+      // NOTE: LongHashedRelation cannot support canMarkRowLookedUp as it cannot
       // handle NULL key
       LongHashedRelation(input, key, sizeEstimate, mm, isNullAware)
     } else {
-      UnsafeHashedRelation(input, key, sizeEstimate, mm, isNullAware, isLookupAware, value)
+      UnsafeHashedRelation(
+        input, key, sizeEstimate, mm, isNullAware, canMarkRowLookedUp, valueExprs)
     }
   }
 }
@@ -344,12 +346,10 @@ private[joins] object UnsafeHashedRelation {
       sizeEstimate: Int,
       taskMemoryManager: TaskMemoryManager,
       isNullAware: Boolean = false,
-      isLookupAware: Boolean = false,
-      value: Option[Seq[Expression]] = None): HashedRelation = {
-    if (isNullAware && isLookupAware) {
-      throw new SparkException(
-        "isLookupAware and isNullAware cannot be enabled at same time for UnsafeHashedRelation")
-    }
+      canMarkRowLookedUp: Boolean = false,
+      valueExprs: Option[Seq[Expression]] = None): HashedRelation = {
+    require(!(isNullAware && canMarkRowLookedUp),
+      "isNullAware and canMarkRowLookedUp cannot be enabled at same time")
 
     val pageSizeBytes = Option(SparkEnv.get).map(_.memoryManager.pageSizeBytes)
       .getOrElse(new SparkConf().get(BUFFER_PAGESIZE).getOrElse(16L * 1024 * 1024))
@@ -376,11 +376,11 @@ private[joins] object UnsafeHashedRelation {
       }
     }
 
-    if (isLookupAware) {
+    if (canMarkRowLookedUp) {
       // Add one extra boolean value at the end as part of the row,
       // to track the information that whether the corresponding key
       // has been looked up or not. See `ShuffledHashJoin.fullOuterJoin` for example of usage.
-      val valueGenerator = UnsafeProjection.create(value.get :+ Literal(false))
+      val valueGenerator = UnsafeProjection.create(valueExprs.get :+ Literal(false))
       while (input.hasNext) {
         val row = input.next().asInstanceOf[UnsafeRow]
         numFields = row.numFields() + 1

sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledHashJoinExec.scala

@@ -66,7 +66,7 @@ case class ShuffledHashJoinExec(
     val start = System.nanoTime()
     val context = TaskContext.get()
 
-    val (isLookupAware, value) = if (joinType == FullOuter) {
+    val (canMarkRowLookedUp, valueExprs) = if (joinType == FullOuter) {
       (true, Some(BindReferences.bindReferences(buildOutput, buildOutput)))
     } else {
       (false, None)
@@ -75,8 +75,8 @@ case class ShuffledHashJoinExec(
       iter,
       buildBoundKeys,
       taskMemoryManager = context.taskMemoryManager(),
-      isLookupAware = isLookupAware,
-      value = value)
+      canMarkRowLookedUp = canMarkRowLookedUp,
+      valueExprs = valueExprs)
     buildTime += NANOSECONDS.toMillis(System.nanoTime() - start)
     buildDataSize += relation.estimatedSize
     // This relation is usually used until the end of task.
@@ -103,7 +103,7 @@ case class ShuffledHashJoinExec(
    * 2. Process rows from stream side by looking up hash relation,
    *    and mark the matched rows from build side be looked up.
    * 3. Process rows from build side by iterating hash relation,
-   *    and filter out rows from build side being looked up already.
+   *    and filter out rows from build side being matched already.
    */
   private def fullOuterJoin(
       streamIter: Iterator[InternalRow],
@@ -180,15 +180,26 @@ case class ShuffledHashJoinExec(
         }
       }
 
-    // Process build side with filtering out rows looked up already
+    // Process build side with filtering out rows looked up and
+    // passed join condition already
+    val streamNullJoinRow = new JoinedRow
+    val streamNullJoinRowWithBuild = {
+      buildSide match {
+        case BuildLeft =>
+          streamNullJoinRow.withRight(streamNullRow)
+          streamNullJoinRow.withLeft _
+        case BuildRight =>
+          streamNullJoinRow.withLeft(streamNullRow)
+          streamNullJoinRow.withRight _
+      }
+    }
     val buildResultIter = hashedRelation.values().flatMap { brow =>
       val unsafebrow = brow.asInstanceOf[UnsafeRow]
       val isLookup = unsafebrow.getBoolean(unsafebrow.numFields() - 1)
       if (!isLookup) {
         val buildRow = buildRowGenerator(unsafebrow)
-        joinRowWithBuild(buildRow)
-        joinRowWithStream(streamNullRow)
-        Some(joinRow)
+        streamNullJoinRowWithBuild(buildRow)
+        Some(streamNullJoinRow)
       } else {
         None
       }

sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala

@@ -1193,31 +1193,42 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
     val inputDFs = Seq(
       // Test unique join key
       (spark.range(10).selectExpr("id as k1"),
-        spark.range(30).selectExpr("id as k2")),
+        spark.range(30).selectExpr("id as k2"),
+        $"k1" === $"k2"),
       // Test non-unique join key
       (spark.range(10).selectExpr("id % 5 as k1"),
-        spark.range(30).selectExpr("id % 5 as k2")),
+        spark.range(30).selectExpr("id % 5 as k2"),
+        $"k1" === $"k2"),
       // Test string join key
       (spark.range(10).selectExpr("cast(id * 3 as string) as k1"),
-        spark.range(30).selectExpr("cast(id as string) as k2")),
+        spark.range(30).selectExpr("cast(id as string) as k2"),
+        $"k1" === $"k2"),
       // Test build side at right
       (spark.range(30).selectExpr("cast(id / 3 as string) as k1"),
-        spark.range(10).selectExpr("cast(id as string) as k2")),
+        spark.range(10).selectExpr("cast(id as string) as k2"),
+        $"k1" === $"k2"),
       // Test NULL join key
       (spark.range(10).map(i => if (i % 2 == 0) i else null).selectExpr("value as k1"),
-        spark.range(30).map(i => if (i % 4 == 0) i else null).selectExpr("value as k2"))
+        spark.range(30).map(i => if (i % 4 == 0) i else null).selectExpr("value as k2"),
+        $"k1" === $"k2"),
+      // Test multiple join keys
+      (spark.range(10).map(i => if (i % 2 == 0) i else null).selectExpr(
+        "value as k1", "cast(value % 5 as short) as k2", "cast(value * 3 as long) as k3"),
+        spark.range(30).map(i => if (i % 4 == 0) i else null).selectExpr(
+          "value as k4", "cast(value % 5 as short) as k5", "cast(value * 3 as long) as k6"),
+        $"k1" === $"k4" && $"k2" === $"k5" && $"k3" === $"k6")
     )
-    inputDFs.foreach { case (df1, df2) =>
+    inputDFs.foreach { case (df1, df2, joinExprs) =>
       withSQLConf(
         SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80",
         SQLConf.SHUFFLE_PARTITIONS.key -> "2") {
-        val smjDF = df1.join(df2, $"k1" === $"k2", "full")
+        val smjDF = df1.join(df2, joinExprs, "full")
         assert(smjDF.queryExecution.executedPlan.collect {
           case _: SortMergeJoinExec => true }.size === 1)
         val smjResult = smjDF.collect()
 
         withSQLConf(SQLConf.PREFER_SORTMERGEJOIN.key -> "false") {
-          val shjDF = df1.join(df2, $"k1" === $"k2", "full")
+          val shjDF = df1.join(df2, joinExprs, "full")
           assert(shjDF.queryExecution.executedPlan.collect {
             case _: ShuffledHashJoinExec => true }.size === 1)
           // Same result between shuffled hash join and sort merge join

sql/core/src/test/scala/org/apache/spark/sql/execution/joins/HashedRelationSuite.scala

@@ -586,22 +586,23 @@ class HashedRelationSuite extends SharedSparkSession {
     val value = Seq(BoundReference(0, IntegerType, true))
     val unsafeProj = UnsafeProjection.create(value)
     val rows = (0 until 100).map(i => unsafeProj(InternalRow(i + 1)).copy())
+    val expectedValues = (0 until 100).map(i => i + 1)
 
     // test LongHashedRelation
     val longRelation = LongHashedRelation(rows.iterator, key, 10, mm)
     var values = longRelation.values()
-    assert(values.map(_.getInt(0)).toArray.sortWith(_ < _) === (0 until 100).map(i => i + 1))
+    assert(values.map(_.getInt(0)).toArray.sortWith(_ < _) === expectedValues)
 
     // test UnsafeHashedRelation
     val unsafeRelation = UnsafeHashedRelation(rows.iterator, key, 10, mm)
     values = unsafeRelation.values()
-    assert(values.map(_.getInt(0)).toArray.sortWith(_ < _) === (0 until 100).map(i => i + 1))
+    assert(values.map(_.getInt(0)).toArray.sortWith(_ < _) === expectedValues)
 
-    // test lookup-aware UnsafeHashedRelation
-    val lookupAwareUnsafeRelation = UnsafeHashedRelation(
-      rows.iterator, key, 10, mm, isLookupAware = true, value = Some(value))
-    values = lookupAwareUnsafeRelation.values()
+    // test UnsafeHashedRelation which can mark row looked up
+    val markRowUnsafeRelation = UnsafeHashedRelation(
+      rows.iterator, key, 10, mm, canMarkRowLookedUp = true, valueExprs = Some(value))
+    values = markRowUnsafeRelation.values()
     assert(values.map(v => (v.getInt(0), v.getBoolean(1))).toArray.sortWith(_._1 < _._1)
-      === (0 until 100).map(i => (i + 1, false)))
+      === expectedValues.map(i => (i, false)))
   }
 }