diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/nullExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/nullExpressions.scala index e7be588c4b465..1d50d8987ccee 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/nullExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/nullExpressions.scala @@ -93,11 +93,37 @@ case class Coalesce(children: Seq[Expression]) } override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { + val resultType = CodeGenerator.javaType(dataType) + // Fast path for the common `coalesce(a, b)` where `b` is non-nullable and its evaluation + // generates no code (a literal or an already-evaluated value), e.g. the `coalesce(sum, 0)` + // emitted for every SUM/AVG update. The result can never be null and `b` carries a pure + // Java expression, so a single ternary replaces the general do-while block below and no + // global mutable isNull state is needed. The empty-code requirement also preserves lazy + // evaluation: there are no statements of `b` to hoist before the null check. + val probedEvals: Option[Seq[ExprCode]] = + if (children.length == 2 && !children(1).nullable) { + val first = children.head.genCode(ctx) + val second = children(1).genCode(ctx) + if (second.code.isEmpty) { + return ev.copy( + code = code""" + |${first.code} + |$resultType ${ev.value} = ${first.isNull} ? ${second.value} : ${first.value}; + """.stripMargin, + isNull = FalseLiteral) + } + // The fast path does not apply, so the general path below reuses the code generated for + // the probe: generating a child twice would duplicate codegen side effects on the context + // (orphaned mutable state for stateful expressions, repeated reference registrations). + Some(Seq(first, second)) + } else { + None + } + ev.isNull = JavaCode.isNullGlobal(ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, ev.isNull)) // all the evals are meant to be in a do { ... } while (false); loop - val evals = children.map { e => - val eval = e.genCode(ctx) + val evals = probedEvals.getOrElse(children.map(_.genCode(ctx))).map { eval => s""" |${eval.code} |if (!${eval.isNull}) { @@ -108,7 +134,6 @@ case class Coalesce(children: Seq[Expression]) """.stripMargin } - val resultType = CodeGenerator.javaType(dataType) val codes = ctx.splitExpressionsWithCurrentInputs( expressions = evals, funcName = "coalesce", diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/NullExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/NullExpressionsSuite.scala index 5c19e69cdfa3d..3d957b662d5da 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/NullExpressionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/NullExpressionsSuite.scala @@ -21,6 +21,7 @@ import java.sql.Timestamp import org.apache.spark.{SparkFunSuite, SparkRuntimeException} import org.apache.spark.sql.catalyst.FunctionIdentifier +import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.analysis.{FunctionRegistry, SimpleAnalyzer, UnresolvedAttribute} import org.apache.spark.sql.catalyst.expressions.codegen.CodegenContext import org.apache.spark.sql.catalyst.expressions.objects.AssertNotNull @@ -207,10 +208,41 @@ class NullExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { test("SPARK-22705: Coalesce should use less global variables") { val ctx = new CodegenContext() - Coalesce(Seq(Literal("a"), Literal("b"))).genCode(ctx) + // Use three children so that the general do-while codegen path is exercised: the two-argument + // form with a non-nullable pure fallback is now special-cased to a ternary with no global + // state (covered by the test below). + Coalesce(Seq(Literal("a"), Literal("b"), Literal("c"))).genCode(ctx) assert(ctx.inlinedMutableStates.size == 1) } + test("Coalesce with a non-nullable pure fallback compiles to a ternary with no global state") { + val ctx = new CodegenContext() + Coalesce(Seq(Literal.create(null, StringType), Literal("b"))).genCode(ctx) + assert(ctx.inlinedMutableStates.isEmpty) + checkEvaluation(Coalesce(Seq(Literal.create(null, StringType), Literal("a"))), "a") + checkEvaluation(Coalesce(Seq(Literal("x"), Literal("y"))), "x") + } + + test("Coalesce ternary fast path over a BoundReference first arg and a primitive fallback") { + // The motivating hot path is coalesce(nullable_column, literal): the first arg is a + // BoundReference read from the input row (a distinct codegen path from a Literal), and the + // fallback is a non-nullable primitive that emits no code, so the ternary fast path applies. + val coalesce = Coalesce(Seq(BoundReference(0, IntegerType, nullable = true), Literal(0))) + checkEvaluation(coalesce, 42, InternalRow(42)) + checkEvaluation(coalesce, 0, InternalRow(null)) + } + + test("Coalesce reuses the probed ExprCode when the non-nullable fallback emits code") { + // A two-arg coalesce whose non-nullable fallback is itself a BoundReference: its code is + // non-empty, so the ternary fast path does not apply and the probed ExprCodes are threaded + // into the general do-while path (rather than generating the children a second time). + val coalesce = Coalesce(Seq( + BoundReference(0, IntegerType, nullable = true), + BoundReference(1, IntegerType, nullable = false))) + checkEvaluation(coalesce, 42, InternalRow(42, 7)) + checkEvaluation(coalesce, 7, InternalRow(null, 7)) + } + test("AtLeastNNonNulls should not throw 64KiB exception") { val inputs = (1 to 4000).map(x => Literal(s"x_$x")) checkEvaluation(AtLeastNNonNulls(1, inputs), true)