From 62a826f17c549ed93300bdce562db56bddd5d959 Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@gmail.com>
Date: Fri, 24 Nov 2017 11:46:58 +0100
Subject: [PATCH] [SPARK-22591][SQL] GenerateOrdering shouldn't change
CodegenContext.INPUT_ROW
## 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.
---
.../expressions/codegen/GenerateOrdering.scala | 16 ++++++++++------
.../sql/catalyst/expressions/OrderingSuite.scala | 11 ++++++++++-
2 files changed, 20 insertions(+), 7 deletions(-)
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
index 1639d1b9dd..4a459571ed 100644
--- 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
index aa61ba2bff..d0604b8eb7 100644
--- 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)
+ }
}
--
GitLab