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 1639d1b9dda1f1eedc6799add395e442c3e33be0..4a459571ed634ece4e8213403b44b22a81bf2fd3 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 aa61ba2bff2bb265961232cbeba8d52501728517..d0604b8eb7675b47f4080016a2123cea1df0de59 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) + } }