Address all comments

c21 · c21 · commit d4e0084f904e · 2020-08-07T10:37:46.000-07:00
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashedRelation.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashedRelation.scala
@@ -96,6 +96,9 @@ 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
    */
   def apply(
       input: Iterator[InternalRow],
@@ -118,6 +121,8 @@ 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
+      // handle NULL key
       LongHashedRelation(input, key, sizeEstimate, mm, isNullAware)
     } else {
       UnsafeHashedRelation(input, key, sizeEstimate, mm, isNullAware, isLookupAware, value)
@@ -148,7 +153,7 @@ private[joins] class UnsafeHashedRelation(
 
   override def estimatedSize: Long = binaryMap.getTotalMemoryConsumption
 
-  // re-used in get()/getValue()
+  // re-used in get()/getValue()/values()
   var resultRow = new UnsafeRow(numFields)
 
   override def get(key: InternalRow): Iterator[InternalRow] = {
@@ -186,6 +191,23 @@ private[joins] class UnsafeHashedRelation(
     }
   }
 
+  override def values(): Iterator[InternalRow] = {
+    val iter = binaryMap.iterator()
+
+    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 keys(): Iterator[InternalRow] = {
     val iter = binaryMap.iterator()
 
@@ -312,23 +334,6 @@ private[joins] class UnsafeHashedRelation(
   override def read(kryo: Kryo, in: Input): Unit = Utils.tryOrIOException {
     read(() => in.readInt(), () => in.readLong(), in.readBytes)
   }
-
-  override def values(): Iterator[InternalRow] = {
-    val iter = binaryMap.iterator()
-
-    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
-      }
-    }
-  }
 }
 
 private[joins] object UnsafeHashedRelation {
@@ -341,6 +346,10 @@ private[joins] object UnsafeHashedRelation {
       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")
+    }
 
     val pageSizeBytes = Option(SparkEnv.get).map(_.memoryManager.pageSizeBytes)
       .getOrElse(new SparkConf().get(BUFFER_PAGESIZE).getOrElse(16L * 1024 * 1024))
@@ -354,44 +363,36 @@ private[joins] object UnsafeHashedRelation {
     val keyGenerator = UnsafeProjection.create(key)
     var numFields = 0
 
+    val append = (key: UnsafeRow, value: UnsafeRow) => {
+      val loc = binaryMap.lookup(key.getBaseObject, key.getBaseOffset, key.getSizeInBytes)
+      val success = loc.append(
+        key.getBaseObject, key.getBaseOffset, key.getSizeInBytes,
+        value.getBaseObject, value.getBaseOffset, value.getSizeInBytes)
+      if (!success) {
+        binaryMap.free()
+        // scalastyle:off throwerror
+        throw new SparkOutOfMemoryError("There is not enough memory to build hash map")
+        // scalastyle:on throwerror
+      }
+    }
+
     if (isLookupAware) {
       // 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))
-
       while (input.hasNext) {
         val row = input.next().asInstanceOf[UnsafeRow]
         numFields = row.numFields() + 1
-        val key = keyGenerator(row)
-        val value = valueGenerator(row)
-        val loc = binaryMap.lookup(key.getBaseObject, key.getBaseOffset, key.getSizeInBytes)
-        val success = loc.append(
-          key.getBaseObject, key.getBaseOffset, key.getSizeInBytes,
-          value.getBaseObject, value.getBaseOffset, value.getSizeInBytes)
-        if (!success) {
-          binaryMap.free()
-          // scalastyle:off throwerror
-          throw new SparkOutOfMemoryError("There is not enough memory to build hash map")
-          // scalastyle:on throwerror
-        }
+        append(keyGenerator(row), valueGenerator(row))
       }
     } else {
       while (input.hasNext) {
         val row = input.next().asInstanceOf[UnsafeRow]
         numFields = row.numFields()
         val key = keyGenerator(row)
         if (!key.anyNull) {
-          val loc = binaryMap.lookup(key.getBaseObject, key.getBaseOffset, key.getSizeInBytes)
-          val success = loc.append(
-            key.getBaseObject, key.getBaseOffset, key.getSizeInBytes,
-            row.getBaseObject, row.getBaseOffset, row.getSizeInBytes)
-          if (!success) {
-            binaryMap.free()
-            // scalastyle:off throwerror
-            throw new SparkOutOfMemoryError("There is not enough memory to build hash map")
-            // scalastyle:on throwerror
-          }
+          append(key, row)
         } else if (isNullAware) {
           return EmptyHashedRelationWithAllNullKeys
         }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledHashJoinExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledHashJoinExec.scala
@@ -66,12 +66,11 @@ case class ShuffledHashJoinExec(
     val start = System.nanoTime()
     val context = TaskContext.get()
 
-    val (isLookupAware, value) =
-      if (joinType == FullOuter) {
-        (true, Some(BindReferences.bindReferences(buildOutput, buildOutput)))
-      } else {
-        (false, None)
-      }
+    val (isLookupAware, value) = if (joinType == FullOuter) {
+      (true, Some(BindReferences.bindReferences(buildOutput, buildOutput)))
+    } else {
+      (false, None)
+    }
     val relation = HashedRelation(
       iter,
       buildBoundKeys,
@@ -110,24 +109,12 @@ case class ShuffledHashJoinExec(
       streamIter: Iterator[InternalRow],
       hashedRelation: HashedRelation,
       numOutputRows: SQLMetric): Iterator[InternalRow] = {
-    abstract class HashJoinedRow extends JoinedRow {
-      /** Updates this JoinedRow by updating its stream side row. Returns itself. */
-      def withStream(newStream: InternalRow): JoinedRow
-
-      /** Updates this JoinedRow by updating its build side row. Returns itself. */
-      def withBuild(newBuild: InternalRow): JoinedRow
-    }
-    val joinRow: HashJoinedRow = buildSide match {
-      case BuildLeft =>
-        new HashJoinedRow {
-          override def withStream(newStream: InternalRow): JoinedRow = withRight(newStream)
-          override def withBuild(newBuild: InternalRow): JoinedRow = withLeft(newBuild)
-        }
-      case BuildRight =>
-        new HashJoinedRow {
-          override def withStream(newStream: InternalRow): JoinedRow = withLeft(newStream)
-          override def withBuild(newBuild: InternalRow): JoinedRow = withRight(newBuild)
-        }
+    val joinRow = new JoinedRow
+    val (joinRowWithStream, joinRowWithBuild) = {
+      buildSide match {
+        case BuildLeft => (joinRow.withRight _, joinRow.withLeft _)
+        case BuildRight => (joinRow.withLeft _, joinRow.withRight _)
+      }
     }
     val joinKeys = streamSideKeyGenerator()
     val buildRowGenerator = UnsafeProjection.create(buildOutput, buildOutput)
@@ -141,31 +128,31 @@ case class ShuffledHashJoinExec(
     val streamResultIter =
       if (hashedRelation.keyIsUnique) {
         streamIter.map { srow =>
-          joinRow.withStream(srow)
+          joinRowWithStream(srow)
           val keys = joinKeys(srow)
           if (keys.anyNull) {
-            joinRow.withBuild(buildNullRow)
+            joinRowWithBuild(buildNullRow)
           } else {
             val matched = hashedRelation.getValue(keys)
             if (matched != null) {
               val buildRow = buildRowGenerator(matched)
-              if (boundCondition(joinRow.withBuild(buildRow))) {
+              if (boundCondition(joinRowWithBuild(buildRow))) {
                 markRowLookedUp(matched.asInstanceOf[UnsafeRow])
                 joinRow
               } else {
-                joinRow.withBuild(buildNullRow)
+                joinRowWithBuild(buildNullRow)
               }
             } else {
-              joinRow.withBuild(buildNullRow)
+              joinRowWithBuild(buildNullRow)
             }
           }
         }
       } else {
         streamIter.flatMap { srow =>
-          joinRow.withStream(srow)
+          joinRowWithStream(srow)
           val keys = joinKeys(srow)
           if (keys.anyNull) {
-            Iterator.single(joinRow.withBuild(buildNullRow))
+            Iterator.single(joinRowWithBuild(buildNullRow))
           } else {
             val buildIter = hashedRelation.get(keys)
             new RowIterator {
@@ -174,14 +161,14 @@ case class ShuffledHashJoinExec(
                 while (buildIter != null && buildIter.hasNext) {
                   val matched = buildIter.next()
                   val buildRow = buildRowGenerator(matched)
-                  if (boundCondition(joinRow.withBuild(buildRow))) {
+                  if (boundCondition(joinRowWithBuild(buildRow))) {
                     markRowLookedUp(matched.asInstanceOf[UnsafeRow])
                     found = true
                     return true
                   }
                 }
                 if (!found) {
-                  joinRow.withBuild(buildNullRow)
+                  joinRowWithBuild(buildNullRow)
                   found = true
                   return true
                 }
@@ -199,8 +186,8 @@ case class ShuffledHashJoinExec(
       val isLookup = unsafebrow.getBoolean(unsafebrow.numFields() - 1)
       if (!isLookup) {
         val buildRow = buildRowGenerator(unsafebrow)
-        joinRow.withBuild(buildRow)
-        joinRow.withStream(streamNullRow)
+        joinRowWithBuild(buildRow)
+        joinRowWithStream(streamNullRow)
         Some(joinRow)
       } else {
         None
@@ -214,7 +201,7 @@ case class ShuffledHashJoinExec(
     }
   }
 
-  // TODO: support full outer shuffled hash join code-gen
+  // TODO(SPARK-32567): support full outer shuffled hash join code-gen
   override def supportCodegen: Boolean = {
     joinType != FullOuter
   }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledJoin.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledJoin.scala
@@ -58,7 +58,7 @@ trait ShuffledJoin extends BaseJoinExec {
         left.output
       case x =>
         throw new IllegalArgumentException(
-          s"ShuffledJoin not take $x as the JoinType")
+          s"${getClass.getSimpleName} not take $x as the JoinType")
     }
   }
 }

Original file line number	Diff line number	Diff line change
`@@ -58,7 +58,7 @@ trait ShuffledJoin extends BaseJoinExec {`
`58`	`58`	`left.output`
`59`	`59`	`case x =>`
`60`	`60`	`throw new IllegalArgumentException(`
`61`		`- s"ShuffledJoin not take $x as the JoinType")`
	`61`	`+ s"${getClass.getSimpleName} not take $x as the JoinType")`
`62`	`62`	`}`
`63`	`63`	`}`
`64`	`64`	`}`