Github user mgaido91 commented on a diff in the pull request: https://github.com/apache/spark/pull/19811#discussion_r156347570 --- Diff: sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala --- @@ -177,11 +189,64 @@ class CodegenContext { * the list of default imports available. * Also, generic type arguments are accepted but ignored. * @param variableName Name of the field. - * @param initCode The statement(s) to put into the init() method to initialize this field. + * @param codeFunctions Function includes statement(s) to put into the init() method to + * initialize this field. An argument is the name of the mutable state variable * If left blank, the field will be default-initialized. + * @param inline whether the declaration and initialization code may be inlined rather than + * compacted. + * @param useFreshName If false and inline is true, the name is not changed + * @return the name of the mutable state variable, which is either the original name if the + * variable is inlined to the outer class, or an array access if the variable is to be + * stored in an array of variables of the same type and initialization. + * There are two use cases. One is to use the original name for global variable instead + * of fresh name. Second is to use the original initialization statement since it is + * complex (e.g. allocate multi-dimensional array or object constructor has varibles). + * Primitive type variables will be inlined into outer class when the total number of + * mutable variables is less than `CodeGenerator.OUTER_CLASS_VARIABLES_THRESHOLD` + * the max size of an array for compaction is given by + * `CodeGenerator.MUTABLESTATEARRAY_SIZE_LIMIT`. */ - def addMutableState(javaType: String, variableName: String, initCode: String = ""): Unit = { - mutableStates += ((javaType, variableName, initCode)) + def addMutableState( + javaType: String, + variableName: String, + codeFunctions: String => String = _ => "", + inline: Boolean = false, + useFreshName: Boolean = true): String = { + val varName = if (useFreshName) freshName(variableName) else variableName + + if (inline || + // want to put a primitive type variable at outerClass for performance + isPrimitiveType(javaType) && + (mutableStates.length < CodeGenerator.OUTER_CLASS_VARIABLES_THRESHOLD) || + // type is multi-dimensional array + javaType.contains("[][]")) { + val initCode = codeFunctions(varName) + mutableStates += ((javaType, varName, initCode)) + varName + } else { + val arrayName = mutableStateArrayCurrentNames.getOrElse(javaType, "") + val prevIdx = mutableStateArrayIdx.getOrElse((javaType, arrayName), -1) + if (0 <= prevIdx && prevIdx < CodeGenerator.MUTABLESTATEARRAY_SIZE_LIMIT - 1) { + // a mutableStateArray for the given type and name has already been declared, + // update the max index of the array and return an array element + val idx = prevIdx + 1 + mutableStateArrayIdx.update((javaType, arrayName), idx) + val initCode = codeFunctions(s"$arrayName[$idx]") + mutableStateArrayInitCodes += initCode + s"$arrayName[$idx]" + } else { + // mutableStateArray has not been declared yet for the given type and name. + // Create a new name for the array, and add an entry to keep track of current array name + // for type. In addition, init code is stored for code generation + val newArrayName = freshName("mutableStateArray") + mutableStateArrayCurrentNames += javaType -> newArrayName + val idx = 0 + mutableStateArrayIdx += (javaType, newArrayName) -> idx --- End diff -- ditto, so we might move this out of the if with the change I suggested few lines above
--- --------------------------------------------------------------------- To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org