Github user viirya commented on a diff in the pull request:
https://github.com/apache/spark/pull/19811#discussion_r154015991
--- Diff:
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
---
@@ -203,18 +271,47 @@ class CodegenContext {
def declareMutableStates(): String = {
// It's possible that we add same mutable state twice, e.g. the
`mergeExpressions` in
// `TypedAggregateExpression`, we should call `distinct` here to
remove the duplicated ones.
- mutableStates.distinct.map { case (javaType, variableName, _) =>
+ val inlinedStates = mutableStates.distinct.map { case (javaType,
variableName, _) =>
s"private $javaType $variableName;"
- }.mkString("\n")
+ }
+
+ val arrayStates = mutableStateArrayInitCodes.map { case (javaType,
arrayName, _) =>
+ val length = mutableStateArrayIdx((javaType, arrayName)) + 1
+ if (javaType.matches("^.*\\[\\]$")) {
+ // initializer had an one-dimensional array variable
+ val baseType = javaType.substring(0, javaType.length - 2)
+ s"private $javaType[] $arrayName = new $baseType[$length][];"
+ } else {
+ // initializer had a scalar variable
+ s"private $javaType[] $arrayName = new $javaType[$length];"
+ }
+ }
+
+ (inlinedStates ++ arrayStates).mkString("\n")
}
def initMutableStates(): String = {
// It's possible that we add same mutable state twice, e.g. the
`mergeExpressions` in
// `TypedAggregateExpression`, we should call `distinct` here to
remove the duplicated ones.
val initCodes = mutableStates.distinct.map(_._3 + "\n")
+ // array state is initialized in loops
+ val arrayInitCodes = mutableStateArrayInitCodes.map {
+ case (javaType, arrayName, qualifiedInitCode) =>
+ if (qualifiedInitCode == "") {
+ ""
+ } else {
+ val loopIdxVar = CodeGenerator.INIT_LOOP_VARIABLE_NAME
+ s"""
+ for (int $loopIdxVar = 0; $loopIdxVar < $arrayName.length;
$loopIdxVar++) {
+ $qualifiedInitCode
--- End diff --
Oh. I see. You put the variables with same init code in the same array.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
