Skip to content

Commit 62a826f

Browse files
viiryacloud-fan
authored andcommitted
[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 <[email protected]> Closes #19800 from viirya/SPARK-22591.
1 parent c121756 commit 62a826f

File tree

2 files changed

+20
-7
lines changed

2 files changed

+20
-7
lines changed

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateOrdering.scala

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -72,13 +72,15 @@ object GenerateOrdering extends CodeGenerator[Seq[SortOrder], Ordering[InternalR
7272
* Generates the code for ordering based on the given order.
7373
*/
7474
def genComparisons(ctx: CodegenContext, ordering: Seq[SortOrder]): String = {
75+
val oldInputRow = ctx.INPUT_ROW
76+
val oldCurrentVars = ctx.currentVars
77+
val inputRow = "i"
78+
ctx.INPUT_ROW = inputRow
79+
// to use INPUT_ROW we must make sure currentVars is null
80+
ctx.currentVars = null
81+
7582
val comparisons = ordering.map { order =>
76-
val oldCurrentVars = ctx.currentVars
77-
ctx.INPUT_ROW = "i"
78-
// to use INPUT_ROW we must make sure currentVars is null
79-
ctx.currentVars = null
8083
val eval = order.child.genCode(ctx)
81-
ctx.currentVars = oldCurrentVars
8284
val asc = order.isAscending
8385
val isNullA = ctx.freshName("isNullA")
8486
val primitiveA = ctx.freshName("primitiveA")
@@ -147,10 +149,12 @@ object GenerateOrdering extends CodeGenerator[Seq[SortOrder], Ordering[InternalR
147149
"""
148150
}.mkString
149151
})
152+
ctx.currentVars = oldCurrentVars
153+
ctx.INPUT_ROW = oldInputRow
150154
// make sure INPUT_ROW is declared even if splitExpressions
151155
// returns an inlined block
152156
s"""
153-
|InternalRow ${ctx.INPUT_ROW} = null;
157+
|InternalRow $inputRow = null;
154158
|$code
155159
""".stripMargin
156160
}

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/OrderingSuite.scala

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ import org.apache.spark.serializer.KryoSerializer
2424
import org.apache.spark.sql.{RandomDataGenerator, Row}
2525
import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow}
2626
import org.apache.spark.sql.catalyst.dsl.expressions._
27-
import org.apache.spark.sql.catalyst.expressions.codegen.{GenerateOrdering, LazilyGeneratedOrdering}
27+
import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, GenerateOrdering, LazilyGeneratedOrdering}
2828
import org.apache.spark.sql.types._
2929

3030
class OrderingSuite extends SparkFunSuite with ExpressionEvalHelper {
@@ -156,4 +156,13 @@ class OrderingSuite extends SparkFunSuite with ExpressionEvalHelper {
156156
assert(genOrdering.compare(rowB1, rowB2) < 0)
157157
}
158158
}
159+
160+
test("SPARK-22591: GenerateOrdering shouldn't change ctx.INPUT_ROW") {
161+
val ctx = new CodegenContext()
162+
ctx.INPUT_ROW = null
163+
164+
val schema = new StructType().add("field", FloatType, nullable = true)
165+
GenerateOrdering.genComparisons(ctx, schema)
166+
assert(ctx.INPUT_ROW == null)
167+
}
159168
}

0 commit comments

Comments
 (0)