Skip to content

Commit f52ea2c

Browse files
committed
: Ported latest fix apache#18972
1 parent 48ea442 commit f52ea2c

File tree

3 files changed

+106
-32
lines changed

3 files changed

+106
-32
lines changed

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -695,6 +695,29 @@ class CodegenContext {
695695
}
696696
}
697697

698+
def createAndAddFunction(
699+
ctx: CodegenContext,
700+
ev: ExprCode,
701+
dataType: DataType,
702+
baseFuncName: String): (String, String, String) = {
703+
val globalIsNull = ctx.freshName("isNull")
704+
ctx.addMutableState("boolean", globalIsNull, s"$globalIsNull = false;")
705+
val globalValue = ctx.freshName("value")
706+
ctx.addMutableState(ctx.javaType(dataType), globalValue,
707+
s"$globalValue = ${ctx.defaultValue(dataType)};")
708+
val funcName = ctx.freshName(baseFuncName)
709+
val funcBody =
710+
s"""
711+
|private void $funcName(InternalRow ${ctx.INPUT_ROW}) {
712+
| ${ev.code.trim}
713+
| $globalIsNull = ${ev.isNull};
714+
| $globalValue = ${ev.value};
715+
|}
716+
""".stripMargin
717+
ctx.addNewFunction(funcName, funcBody)
718+
(funcName, globalIsNull, globalValue)
719+
}
720+
698721
/**
699722
* Perform a function which generates a sequence of ExprCodes with a given mapping between
700723
* expressions and common expressions, instead of using the mapping in current context.

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

Lines changed: 3 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -72,11 +72,11 @@ case class If(predicate: Expression, trueValue: Expression, falseValue: Expressi
7272
(ctx.INPUT_ROW != null && ctx.currentVars == null)) {
7373

7474
val (condFuncName, condGlobalIsNull, condGlobalValue) =
75-
createAndAddFunction(ctx, condEval, predicate.dataType, "evalIfCondExpr")
75+
ctx.createAndAddFunction(ctx, condEval, predicate.dataType, "evalIfCondExpr")
7676
val (trueFuncName, trueGlobalIsNull, trueGlobalValue) =
77-
createAndAddFunction(ctx, trueEval, trueValue.dataType, "evalIfTrueExpr")
77+
ctx.createAndAddFunction(ctx, trueEval, trueValue.dataType, "evalIfTrueExpr")
7878
val (falseFuncName, falseGlobalIsNull, falseGlobalValue) =
79-
createAndAddFunction(ctx, falseEval, falseValue.dataType, "evalIfFalseExpr")
79+
ctx.createAndAddFunction(ctx, falseEval, falseValue.dataType, "evalIfFalseExpr")
8080
s"""
8181
$condFuncName(${ctx.INPUT_ROW});
8282
boolean ${ev.isNull} = false;
@@ -112,29 +112,6 @@ case class If(predicate: Expression, trueValue: Expression, falseValue: Expressi
112112
ev.copy(code = generatedCode)
113113
}
114114

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)
136-
}
137-
138115
override def toString: String = s"if ($predicate) $trueValue else $falseValue"
139116

140117
override def sql: String = s"(IF(${predicate.sql}, ${trueValue.sql}, ${falseValue.sql}))"

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

Lines changed: 80 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -293,11 +293,49 @@ case class And(left: Expression, right: Expression) extends BinaryOperator with
293293
val eval2 = right.genCode(ctx)
294294

295295
// The result should be `false`, if any of them is `false` whenever the other is null or not.
296-
if (!left.nullable && !right.nullable) {
296+
297+
// place generated code of eval1 and eval2 in separate methods if their code combined is large
298+
val combinedLength = eval1.code.length + eval2.code.length
299+
if (combinedLength > 1024 &&
300+
// Split these expressions only if they are created from a row object
301+
(ctx.INPUT_ROW != null && ctx.currentVars == null)) {
302+
303+
val (eval1FuncName, eval1GlobalIsNull, eval1GlobalValue) =
304+
ctx.createAndAddFunction(ctx, eval1, BooleanType, "eval1Expr")
305+
val (eval2FuncName, eval2GlobalIsNull, eval2GlobalValue) =
306+
ctx.createAndAddFunction(ctx, eval2, BooleanType, "eval2Expr")
307+
if (!left.nullable && !right.nullable) {
308+
val generatedCode = s"""
309+
$eval1FuncName(${ctx.INPUT_ROW});
310+
boolean ${ev.value} = false;
311+
if (${eval1GlobalValue}) {
312+
$eval2FuncName(${ctx.INPUT_ROW});
313+
${ev.value} = ${eval2GlobalValue};
314+
}
315+
"""
316+
ev.copy(code = generatedCode, isNull = "false")
317+
} else {
318+
val generatedCode = s"""
319+
$eval1FuncName(${ctx.INPUT_ROW});
320+
boolean ${ev.isNull} = false;
321+
boolean ${ev.value} = false;
322+
if (!${eval1GlobalIsNull} && !${eval1GlobalValue}) {
323+
} else {
324+
$eval2FuncName(${ctx.INPUT_ROW});
325+
if (!${eval2GlobalIsNull} && !${eval2GlobalValue}) {
326+
} else if (!${eval1GlobalIsNull} && !${eval2GlobalIsNull}) {
327+
${ev.value} = true;
328+
} else {
329+
${ev.isNull} = true;
330+
}
331+
}
332+
"""
333+
ev.copy(code = generatedCode)
334+
}
335+
} else if (!left.nullable && !right.nullable) {
297336
ev.copy(code = s"""
298337
${eval1.code}
299338
boolean ${ev.value} = false;
300-
301339
if (${eval1.value}) {
302340
${eval2.code}
303341
${ev.value} = ${eval2.value};
@@ -307,7 +345,6 @@ case class And(left: Expression, right: Expression) extends BinaryOperator with
307345
${eval1.code}
308346
boolean ${ev.isNull} = false;
309347
boolean ${ev.value} = false;
310-
311348
if (!${eval1.isNull} && !${eval1.value}) {
312349
} else {
313350
${eval2.code}
@@ -356,12 +393,50 @@ case class Or(left: Expression, right: Expression) extends BinaryOperator with P
356393
val eval2 = right.genCode(ctx)
357394

358395
// The result should be `true`, if any of them is `true` whenever the other is null or not.
359-
if (!left.nullable && !right.nullable) {
396+
397+
// place generated code of eval1 and eval2 in separate methods if their code combined is large
398+
val combinedLength = eval1.code.length + eval2.code.length
399+
if (combinedLength > 1024 &&
400+
// Split these expressions only if they are created from a row object
401+
(ctx.INPUT_ROW != null && ctx.currentVars == null)) {
402+
403+
val (eval1FuncName, eval1GlobalIsNull, eval1GlobalValue) =
404+
ctx.createAndAddFunction(ctx, eval1, BooleanType, "eval1Expr")
405+
val (eval2FuncName, eval2GlobalIsNull, eval2GlobalValue) =
406+
ctx.createAndAddFunction(ctx, eval2, BooleanType, "eval2Expr")
407+
if (!left.nullable && !right.nullable) {
408+
val generatedCode = s"""
409+
$eval1FuncName(${ctx.INPUT_ROW});
410+
boolean ${ev.value} = true;
411+
if (!${eval1GlobalValue}) {
412+
$eval2FuncName(${ctx.INPUT_ROW});
413+
${ev.value} = ${eval2GlobalValue};
414+
}
415+
"""
416+
ev.copy(code = generatedCode, isNull = "false")
417+
} else {
418+
val generatedCode = s"""
419+
$eval1FuncName(${ctx.INPUT_ROW});
420+
boolean ${ev.isNull} = false;
421+
boolean ${ev.value} = true;
422+
if (!${eval1GlobalIsNull} && ${eval1GlobalValue}) {
423+
} else {
424+
$eval2FuncName(${ctx.INPUT_ROW});
425+
if (!${eval2GlobalIsNull} && ${eval2GlobalValue}) {
426+
} else if (!${eval1GlobalIsNull} && !${eval2GlobalIsNull}) {
427+
${ev.value} = false;
428+
} else {
429+
${ev.isNull} = true;
430+
}
431+
}
432+
"""
433+
ev.copy(code = generatedCode)
434+
}
435+
} else if (!left.nullable && !right.nullable) {
360436
ev.isNull = "false"
361437
ev.copy(code = s"""
362438
${eval1.code}
363439
boolean ${ev.value} = true;
364-
365440
if (!${eval1.value}) {
366441
${eval2.code}
367442
${ev.value} = ${eval2.value};
@@ -371,7 +446,6 @@ case class Or(left: Expression, right: Expression) extends BinaryOperator with P
371446
${eval1.code}
372447
boolean ${ev.isNull} = false;
373448
boolean ${ev.value} = true;
374-
375449
if (!${eval1.isNull} && ${eval1.value}) {
376450
} else {
377451
${eval2.code}

0 commit comments

Comments
 (0)