Fix the cases where operators set up its produce framework.

viirya · viirya · commit 1694c9b3429b · 2017-08-14T06:35:27.000Z
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SortExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SortExec.scala
@@ -177,7 +177,7 @@ case class SortExec(
      """.stripMargin.trim
   }
 
-  override protected def effectiveContinueStatement: String = "continue;"
+  override protected def isConsumeInSeparateFunc: Boolean = false
 
   protected override val shouldStopRequired = false
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/WholeStageCodegenExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/WholeStageCodegenExec.scala
@@ -199,6 +199,10 @@ trait CodegenSupport extends SparkPlan {
    *     ... // logic of Op2Exec to consume rows.
    *   }
    * For now, `isConsumeInSeparateFunc` of Op2Exec will be `true`.
+   *
+   * Notice for some operators like `HashAggregateExec`, it doesn't chain previous consume functions
+   * but begins with its produce framework. We should override `isConsumeInSeparateFunc` to return
+   * `false`.
    */
   protected def isConsumeInSeparateFunc: Boolean = {
     val codegenChildren = children.map(_.asInstanceOf[CodegenSupport])
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala
@@ -163,7 +163,7 @@ case class HashAggregateExec(
   // The variables used as aggregation buffer. Only used for aggregation without keys.
   private var bufVars: Seq[ExprCode] = _
 
-  override protected def effectiveContinueStatement: String = "continue;"
+  override protected def isConsumeInSeparateFunc: Boolean = false
 
   private def doProduceWithoutKeys(ctx: CodegenContext): String = {
     val initAgg = ctx.freshName("initAgg")
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoinExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoinExec.scala
@@ -547,6 +547,8 @@ case class SortMergeJoinExec(
     }
   }
 
+  override protected def isConsumeInSeparateFunc: Boolean = false
+
   override def doProduce(ctx: CodegenContext): String = {
     ctx.copyResult = true
     val leftInput = ctx.freshName("leftInput")

Original file line number	Diff line number	Diff line change
`@@ -177,7 +177,7 @@ case class SortExec(`
`177`	`177`	`""".stripMargin.trim`
`178`	`178`	`}`
`179`	`179`
`180`		`- override protected def effectiveContinueStatement: String = "continue;"`
	`180`	`+ override protected def isConsumeInSeparateFunc: Boolean = false`
`181`	`181`
`182`	`182`	`protected override val shouldStopRequired = false`
`183`	`183`
Original file line number	Diff line number	Diff line change
`@@ -547,6 +547,8 @@ case class SortMergeJoinExec(`
`547`	`547`	`}`
`548`	`548`	`}`
`549`	`549`
	`550`	`+ override protected def isConsumeInSeparateFunc: Boolean = false`
	`551`	`+`
`550`	`552`	`override def doProduce(ctx: CodegenContext): String = {`
`551`	`553`	`ctx.copyResult = true`
`552`	`554`	`val leftInput = ctx.freshName("leftInput")`