Github user cloud-fan commented on a diff in the pull request:
https://github.com/apache/spark/pull/19811#discussion_r156926146
--- Diff:
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CodeGenerationSuite.scala
---
@@ -401,4 +401,26 @@ class CodeGenerationSuite extends SparkFunSuite with
ExpressionEvalHelper {
ctx.addReferenceObj("foo", foo)
assert(ctx.mutableStates.isEmpty)
}
+
+ test("SPARK-18016: def mutable states by using an array") {
+ val ctx1 = new CodegenContext
+ for (i <- 1 to CodeGenerator.OUTER_CLASS_VARIABLES_THRESHOLD + 10) {
+ ctx1.addMutableState(ctx1.JAVA_INT, "i", v => s"$v = $i;")
+ }
+ assert(ctx1.mutableStates.size ==
CodeGenerator.OUTER_CLASS_VARIABLES_THRESHOLD)
+ // When the number of primitive type mutable states is over the
threshold, others are
+ // allocated into an array
+
assert(ctx1.mutableStateArrayMap.get(ctx1.JAVA_INT).get.arrayNames.size == 1)
+ assert(ctx1.mutableStateInitCodes.size ==
CodeGenerator.OUTER_CLASS_VARIABLES_THRESHOLD + 10)
+
+ val ctx2 = new CodegenContext
+ for (i <- 1 to CodeGenerator.MUTABLESTATEARRAY_SIZE_LIMIT + 10) {
+ ctx2.addMutableState("InternalRow[]", "r", v => s"$v = new
InternalRow[$i];")
+ }
+ // When the number of non-primitive type mutable states is over the
threshold, others are
+ // allocated into a new array
+
assert(ctx2.mutableStateArrayMap.get("InternalRow[]").get.arrayNames.size == 2)
--- End diff --
and `assert(ctx2.mutableStateArrayMap("InternalRow[]").getCurrentIndex ==
10)`
---
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
