Skip to content

Commit e463678

Browse files
Kapil Singhcloud-fan
authored andcommitted
[SPARK-18091][SQL] Deep if expressions cause Generated SpecificUnsafeProjection code to exceed JVM code size limit
## What changes were proposed in this pull request? Fix for SPARK-18091 which is a bug related to large if expressions causing generated SpecificUnsafeProjection code to exceed JVM code size limit. This PR changes if expression's code generation to place its predicate, true value and false value expressions' generated code in separate methods in context so as to never generate too long combined code. ## How was this patch tested? Added a unit test and also tested manually with the application (having transformations similar to the unit test) which caused the issue to be identified in the first place. Author: Kapil Singh <[email protected]> Closes #15620 from kapilsingh5050/SPARK-18091-IfCodegenFix.
1 parent edb0ad9 commit e463678

File tree

2 files changed

+90
-13
lines changed

2 files changed

+90
-13
lines changed

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/conditionalExpressions.scala

Lines changed: 69 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -64,19 +64,75 @@ case class If(predicate: Expression, trueValue: Expression, falseValue: Expressi
6464
val trueEval = trueValue.genCode(ctx)
6565
val falseEval = falseValue.genCode(ctx)
6666

67-
ev.copy(code = s"""
68-
${condEval.code}
69-
boolean ${ev.isNull} = false;
70-
${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
71-
if (!${condEval.isNull} && ${condEval.value}) {
72-
${trueEval.code}
73-
${ev.isNull} = ${trueEval.isNull};
74-
${ev.value} = ${trueEval.value};
75-
} else {
76-
${falseEval.code}
77-
${ev.isNull} = ${falseEval.isNull};
78-
${ev.value} = ${falseEval.value};
79-
}""")
67+
// place generated code of condition, true value and false value in separate methods if
68+
// their code combined is large
69+
val combinedLength = condEval.code.length + trueEval.code.length + falseEval.code.length
70+
val generatedCode = if (combinedLength > 1024 &&
71+
// Split these expressions only if they are created from a row object
72+
(ctx.INPUT_ROW != null && ctx.currentVars == null)) {
73+
74+
val (condFuncName, condGlobalIsNull, condGlobalValue) =
75+
createAndAddFunction(ctx, condEval, predicate.dataType, "evalIfCondExpr")
76+
val (trueFuncName, trueGlobalIsNull, trueGlobalValue) =
77+
createAndAddFunction(ctx, trueEval, trueValue.dataType, "evalIfTrueExpr")
78+
val (falseFuncName, falseGlobalIsNull, falseGlobalValue) =
79+
createAndAddFunction(ctx, falseEval, falseValue.dataType, "evalIfFalseExpr")
80+
s"""
81+
$condFuncName(${ctx.INPUT_ROW});
82+
boolean ${ev.isNull} = false;
83+
${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
84+
if (!$condGlobalIsNull && $condGlobalValue) {
85+
$trueFuncName(${ctx.INPUT_ROW});
86+
${ev.isNull} = $trueGlobalIsNull;
87+
${ev.value} = $trueGlobalValue;
88+
} else {
89+
$falseFuncName(${ctx.INPUT_ROW});
90+
${ev.isNull} = $falseGlobalIsNull;
91+
${ev.value} = $falseGlobalValue;
92+
}
93+
"""
94+
}
95+
else {
96+
s"""
97+
${condEval.code}
98+
boolean ${ev.isNull} = false;
99+
${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
100+
if (!${condEval.isNull} && ${condEval.value}) {
101+
${trueEval.code}
102+
${ev.isNull} = ${trueEval.isNull};
103+
${ev.value} = ${trueEval.value};
104+
} else {
105+
${falseEval.code}
106+
${ev.isNull} = ${falseEval.isNull};
107+
${ev.value} = ${falseEval.value};
108+
}
109+
"""
110+
}
111+
112+
ev.copy(code = generatedCode)
113+
}
114+
115+
private def createAndAddFunction(
116+
ctx: CodegenContext,
117+
ev: ExprCode,
118+
dataType: DataType,
119+
baseFuncName: String): (String, String, String) = {
120+
val globalIsNull = ctx.freshName("isNull")
121+
ctx.addMutableState("boolean", globalIsNull, s"$globalIsNull = false;")
122+
val globalValue = ctx.freshName("value")
123+
ctx.addMutableState(ctx.javaType(dataType), globalValue,
124+
s"$globalValue = ${ctx.defaultValue(dataType)};")
125+
val funcName = ctx.freshName(baseFuncName)
126+
val funcBody =
127+
s"""
128+
|private void $funcName(InternalRow ${ctx.INPUT_ROW}) {
129+
| ${ev.code.trim}
130+
| $globalIsNull = ${ev.isNull};
131+
| $globalValue = ${ev.value};
132+
|}
133+
""".stripMargin
134+
ctx.addNewFunction(funcName, funcBody)
135+
(funcName, globalIsNull, globalValue)
80136
}
81137

82138
override def toString: String = s"if ($predicate) $trueValue else $falseValue"

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

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,27 @@ class CodeGenerationSuite extends SparkFunSuite with ExpressionEvalHelper {
9797
assert(actual(0) == cases)
9898
}
9999

100+
test("SPARK-18091: split large if expressions into blocks due to JVM code size limit") {
101+
val inStr = "StringForTesting"
102+
val row = create_row(inStr)
103+
val inputStrAttr = 'a.string.at(0)
104+
105+
var strExpr: Expression = inputStrAttr
106+
for (_ <- 1 to 13) {
107+
strExpr = If(EqualTo(Decode(Encode(strExpr, "utf-8"), "utf-8"), inputStrAttr),
108+
strExpr, strExpr)
109+
}
110+
111+
val expressions = Seq(strExpr)
112+
val plan = GenerateUnsafeProjection.generate(expressions, true)
113+
val actual = plan(row).toSeq(expressions.map(_.dataType))
114+
val expected = Seq(UTF8String.fromString(inStr))
115+
116+
if (!checkResult(actual, expected)) {
117+
fail(s"Incorrect Evaluation: expressions: $expressions, actual: $actual, expected: $expected")
118+
}
119+
}
120+
100121
test("SPARK-14793: split wide array creation into blocks due to JVM code size limit") {
101122
val length = 5000
102123
val expressions = Seq(CreateArray(List.fill(length)(EqualTo(Literal(1), Literal(1)))))

0 commit comments

Comments
 (0)