[SPARK-22591][SQL] GenerateOrdering shouldn't change CodegenContext.INPUT_ROW

viirya · cloud-fan · commit 62a826f17c54 · 2017-11-24T11:46:58.000+01:00
## What changes were proposed in this pull request? When I played with codegen in developing another PR, I found the value of `CodegenContext.INPUT_ROW` is not reliable. Under wholestage codegen, it is assigned to null first and then suddenly changed to `i`. The reason is `GenerateOrdering` changes `CodegenContext.INPUT_ROW` but doesn't restore it back. ## How was this patch tested? Added test. Author: Liang-Chi Hsieh <viirya@gmail.com> Closes #19800 from viirya/SPARK-22591.
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateOrdering.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateOrdering.scala
@@ -72,13 +72,15 @@ object GenerateOrdering extends CodeGenerator[Seq[SortOrder], Ordering[InternalR
    * Generates the code for ordering based on the given order.
    */
   def genComparisons(ctx: CodegenContext, ordering: Seq[SortOrder]): String = {
+    val oldInputRow = ctx.INPUT_ROW
+    val oldCurrentVars = ctx.currentVars
+    val inputRow = "i"
+    ctx.INPUT_ROW = inputRow
+    // to use INPUT_ROW we must make sure currentVars is null
+    ctx.currentVars = null
+
     val comparisons = ordering.map { order =>
-      val oldCurrentVars = ctx.currentVars
-      ctx.INPUT_ROW = "i"
-      // to use INPUT_ROW we must make sure currentVars is null
-      ctx.currentVars = null
       val eval = order.child.genCode(ctx)
-      ctx.currentVars = oldCurrentVars
       val asc = order.isAscending
       val isNullA = ctx.freshName("isNullA")
       val primitiveA = ctx.freshName("primitiveA")
@@ -147,10 +149,12 @@ object GenerateOrdering extends CodeGenerator[Seq[SortOrder], Ordering[InternalR
           """
         }.mkString
       })
+    ctx.currentVars = oldCurrentVars
+    ctx.INPUT_ROW = oldInputRow
     // make sure INPUT_ROW is declared even if splitExpressions
     // returns an inlined block
     s"""
-       |InternalRow ${ctx.INPUT_ROW} = null;
+       |InternalRow $inputRow = null;
        |$code
      """.stripMargin
   }
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/OrderingSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/OrderingSuite.scala
@@ -24,7 +24,7 @@ import org.apache.spark.serializer.KryoSerializer
 import org.apache.spark.sql.{RandomDataGenerator, Row}
 import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow}
 import org.apache.spark.sql.catalyst.dsl.expressions._
-import org.apache.spark.sql.catalyst.expressions.codegen.{GenerateOrdering, LazilyGeneratedOrdering}
+import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, GenerateOrdering, LazilyGeneratedOrdering}
 import org.apache.spark.sql.types._
 
 class OrderingSuite extends SparkFunSuite with ExpressionEvalHelper {
@@ -156,4 +156,13 @@ class OrderingSuite extends SparkFunSuite with ExpressionEvalHelper {
       assert(genOrdering.compare(rowB1, rowB2) < 0)
     }
   }
+
+  test("SPARK-22591: GenerateOrdering shouldn't change ctx.INPUT_ROW") {
+    val ctx = new CodegenContext()
+    ctx.INPUT_ROW = null
+
+    val schema = new StructType().add("field", FloatType, nullable = true)
+    GenerateOrdering.genComparisons(ctx, schema)
+    assert(ctx.INPUT_ROW == null)
+  }
 }
