[SPARK-57902][SQL] Emit a ternary for 2-arg Coalesce with a non-nullable pure fallback

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",

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)