Address comments for prepareRelation and others

Author: c21

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

@@ -400,7 +400,7 @@ trait HashJoin extends BaseJoinExec with CodegenSupport {
    * Generates the code for Inner join.
    */
   protected def codegenInner(ctx: CodegenContext, input: Seq[ExprCode]): String = {
-    val (relationTerm, keyIsKnownUnique) = prepareRelation(ctx)
+    val (relationTerm, keyIsUnique) = prepareRelation(ctx)
     val (keyEv, anyNull) = genStreamSideJoinKey(ctx, input)
     val (matched, checkCondition, buildVars) = getJoinCondition(ctx, input)
     val numOutput = metricTerm(ctx, "numOutputRows")
@@ -410,7 +410,7 @@ trait HashJoin extends BaseJoinExec with CodegenSupport {
       case BuildRight => input ++ buildVars
     }
 
-    if (keyIsKnownUnique) {
+    if (keyIsUnique) {
       s"""
          |// generate join key for stream side
          |${keyEv.code}
@@ -450,7 +450,7 @@ trait HashJoin extends BaseJoinExec with CodegenSupport {
    * Generates the code for left or right outer join.
    */
   protected def codegenOuter(ctx: CodegenContext, input: Seq[ExprCode]): String = {
-    val (relationTerm, keyIsKnownUnique) = prepareRelation(ctx)
+    val (relationTerm, keyIsUnique) = prepareRelation(ctx)
     val (keyEv, anyNull) = genStreamSideJoinKey(ctx, input)
     val matched = ctx.freshName("matched")
     val buildVars = genBuildSideVars(ctx, matched)
@@ -482,7 +482,7 @@ trait HashJoin extends BaseJoinExec with CodegenSupport {
       case BuildRight => input ++ buildVars
     }
 
-    if (keyIsKnownUnique) {
+    if (keyIsUnique) {
       s"""
          |// generate join key for stream side
          |${keyEv.code}
@@ -527,12 +527,12 @@ trait HashJoin extends BaseJoinExec with CodegenSupport {
    * Generates the code for left semi join.
    */
   protected def codegenSemi(ctx: CodegenContext, input: Seq[ExprCode]): String = {
-    val (relationTerm, keyIsKnownUnique) = prepareRelation(ctx)
+    val (relationTerm, keyIsUnique) = prepareRelation(ctx)
     val (keyEv, anyNull) = genStreamSideJoinKey(ctx, input)
     val (matched, checkCondition, _) = getJoinCondition(ctx, input)
     val numOutput = metricTerm(ctx, "numOutputRows")
 
-    if (keyIsKnownUnique) {
+    if (keyIsUnique) {
       s"""
          |// generate join key for stream side
          |${keyEv.code}
@@ -576,12 +576,12 @@ trait HashJoin extends BaseJoinExec with CodegenSupport {
    * Generates the code for anti join.
    */
   protected def codegenAnti(ctx: CodegenContext, input: Seq[ExprCode]): String = {
-    val (relationTerm, keyIsKnownUnique) = prepareRelation(ctx)
+    val (relationTerm, keyIsUnique) = prepareRelation(ctx)
     val (keyEv, anyNull) = genStreamSideJoinKey(ctx, input)
     val (matched, checkCondition, _) = getJoinCondition(ctx, input)
     val numOutput = metricTerm(ctx, "numOutputRows")
 
-    if (keyIsKnownUnique) {
+    if (keyIsUnique) {
       val found = ctx.freshName("found")
       s"""
          |boolean $found = false;
@@ -637,7 +637,7 @@ trait HashJoin extends BaseJoinExec with CodegenSupport {
    * Generates the code for existence join.
    */
   protected def codegenExistence(ctx: CodegenContext, input: Seq[ExprCode]): String = {
-    val (relationTerm, keyIsKnownUnique) = prepareRelation(ctx)
+    val (relationTerm, keyIsUnique) = prepareRelation(ctx)
     val (keyEv, anyNull) = genStreamSideJoinKey(ctx, input)
     val numOutput = metricTerm(ctx, "numOutputRows")
     val existsVar = ctx.freshName("exists")
@@ -664,7 +664,7 @@ trait HashJoin extends BaseJoinExec with CodegenSupport {
     val resultVar = input ++ Seq(ExprCode.forNonNullValue(
       JavaCode.variable(existsVar, BooleanType)))
 
-    if (keyIsKnownUnique) {
+    if (keyIsUnique) {
       s"""
          |// generate join key for stream side
          |${keyEv.code}

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

@@ -85,16 +85,8 @@ case class ShuffledHashJoinExec(
     // inline mutable state since not many join operations in a task
     val streamedInput = ctx.addMutableState(
       "scala.collection.Iterator", "streamedInput", v => s"$v = inputs[0];", forceInline = true)
-    val buildInput = ctx.addMutableState(
-      "scala.collection.Iterator", "buildInput", v => s"$v = inputs[1];", forceInline = true)
-    val initRelation = ctx.addMutableState(
-      CodeGenerator.JAVA_BOOLEAN, "initRelation", v => s"$v = false;", forceInline = true)
     val streamedRow = ctx.addMutableState(
       "InternalRow", "streamedRow", forceInline = true)
-
-    val thisPlan = ctx.addReferenceObj("plan", this)
-    val (relationTerm, _) = prepareRelation(ctx)
-    val buildRelation = s"$relationTerm = $thisPlan.buildHashedRelation($buildInput);"
     val (streamInputVar, streamInputVarDecl) = createVars(ctx, streamedRow, streamedPlan.output)
 
     val join = joinType match {
@@ -109,12 +101,6 @@ case class ShuffledHashJoinExec(
     }
 
     s"""
-       |// construct hash map for shuffled hash join build side
-       |if (!$initRelation) {
-       |  $buildRelation
-       |  $initRelation = true;
-       |}
-       |
        |while ($streamedInput.hasNext()) {
        |  $streamedRow = (InternalRow) $streamedInput.next();
        |  ${streamInputVarDecl.mkString("\n")}
@@ -130,13 +116,12 @@ case class ShuffledHashJoinExec(
    * and boolean false to indicate key not to be known unique in code-gen time.
    */
   protected override def prepareRelation(ctx: CodegenContext): (String, Boolean) = {
-    if (relationTerm == null) {
-      // Inline mutable state since not many join operations in a task
-      relationTerm = ctx.addMutableState(
-        "org.apache.spark.sql.execution.joins.HashedRelation", "relation", forceInline = true)
-    }
+    val thisPlan = ctx.addReferenceObj("plan", this)
+    val clsName = classOf[HashedRelation].getName
+
+    // Inline mutable state since not many join operations in a task
+    val relationTerm = ctx.addMutableState(clsName, "relation",
+      v => s"$v = $thisPlan.buildHashedRelation(inputs[1]);", forceInline = true)
     (relationTerm, false)
   }
-
-  private var relationTerm: String = _
 }