Address all comments and add unit test

Author: c21

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

@@ -135,18 +135,15 @@ private[execution] object HashedRelation {
 private[joins] class UnsafeHashedRelation(
     private var numKeys: Int,
     private var numFields: Int,
-    private var binaryMap: BytesToBytesMap,
-    private val isLookupAware: Boolean = false)
+    private var binaryMap: BytesToBytesMap)
   extends HashedRelation with Externalizable with KryoSerializable {
 
-  private[joins] def this() = this(0, 0, null, false)  // Needed for serialization
+  private[joins] def this() = this(0, 0, null)  // Needed for serialization
 
-  override def keyIsUnique: Boolean = {
-    binaryMap.numKeys() == binaryMap.numValues()
-  }
+  override def keyIsUnique: Boolean = binaryMap.numKeys() == binaryMap.numValues()
 
   override def asReadOnlyCopy(): UnsafeHashedRelation = {
-    new UnsafeHashedRelation(numKeys, numFields, binaryMap, isLookupAware)
+    new UnsafeHashedRelation(numKeys, numFields, binaryMap)
   }
 
   override def estimatedSize: Long = binaryMap.getTotalMemoryConsumption
@@ -317,23 +314,19 @@ private[joins] class UnsafeHashedRelation(
   }
 
   override def values(): Iterator[InternalRow] = {
-    if (isLookupAware) {
-      val iter = binaryMap.iterator()
+    val iter = binaryMap.iterator()
 
-      new Iterator[InternalRow] {
-        override def hasNext: Boolean = iter.hasNext
+    new Iterator[InternalRow] {
+      override def hasNext: Boolean = iter.hasNext
 
-        override def next(): InternalRow = {
-          if (!hasNext) {
-            throw new NoSuchElementException("End of the iterator")
-          }
-          val loc = iter.next()
-          resultRow.pointTo(loc.getValueBase, loc.getValueOffset, loc.getValueLength)
-          resultRow
+      override def next(): InternalRow = {
+        if (!hasNext) {
+          throw new NoSuchElementException("End of the iterator")
         }
+        val loc = iter.next()
+        resultRow.pointTo(loc.getValueBase, loc.getValueOffset, loc.getValueLength)
+        resultRow
       }
-    } else {
-      throw new UnsupportedOperationException
     }
   }
 }
@@ -405,7 +398,7 @@ private[joins] object UnsafeHashedRelation {
       }
     }
 
-    new UnsafeHashedRelation(key.size, numFields, binaryMap, isLookupAware)
+    new UnsafeHashedRelation(key.size, numFields, binaryMap)
   }
 }
 
@@ -945,7 +938,9 @@ class LongHashedRelation(
   override def keys(): Iterator[InternalRow] = map.keys()
 
   override def values(): Iterator[InternalRow] = {
-    throw new UnsupportedOperationException
+    keys().flatMap { key =>
+      get(key.getLong(0))
+    }
   }
 }
 

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

@@ -134,11 +134,8 @@ case class ShuffledHashJoinExec(
     val buildNullRow = new GenericInternalRow(buildOutput.length)
     val streamNullRow = new GenericInternalRow(streamedOutput.length)
 
-    def markRowLookedUp(row: UnsafeRow): Unit = {
-      if (!row.getBoolean(row.numFields() - 1)) {
-        row.setBoolean(row.numFields() - 1, true)
-      }
-    }
+    def markRowLookedUp(row: UnsafeRow): Unit =
+      row.setBoolean(row.numFields() - 1, true)
 
     // Process stream side with looking up hash relation
     val streamResultIter =

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

@@ -1189,7 +1189,7 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
     }
   }
 
-  test("Full outer shuffled hash join") {
+  test("SPARK-32399: Full outer shuffled hash join") {
     val inputDFs = Seq(
       // Test unique join key
       (spark.range(10).selectExpr("id as k1"),

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

@@ -580,4 +580,28 @@ class HashedRelationSuite extends SharedSparkSession {
       assert(proj(packedKeys).get(0, dt) == -i - 1)
     }
   }
+
+  test("SPARK-32399: test values() method for HashedRelation") {
+    val key = Seq(BoundReference(0, LongType, false))
+    val value = Seq(BoundReference(0, IntegerType, true))
+    val unsafeProj = UnsafeProjection.create(value)
+    val rows = (0 until 100).map(i => unsafeProj(InternalRow(i + 1)).copy())
+
+    // 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))
+
+    // 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))
+
+    // test lookup-aware UnsafeHashedRelation
+    val lookupAwareUnsafeRelation = UnsafeHashedRelation(
+      rows.iterator, key, 10, mm, isLookupAware = true, value = Some(value))
+    values = lookupAwareUnsafeRelation.values()
+    assert(values.map(v => (v.getInt(0), v.getBoolean(1))).toArray.sortWith(_._1 < _._1)
+      === (0 until 100).map(i => (i + 1, false)))
+  }
 }