From 8a45067f0ff8806cb6bf030f195eb914c24fbb6f Mon Sep 17 00:00:00 2001 From: Gengliang Wang Date: Thu, 2 Jul 2026 20:36:00 +0000 Subject: [PATCH 1/3] [SPARK-57902][SQL] Emit a ternary for 2-arg Coalesce with a non-nullable pure fallback ### What changes were proposed in this pull request? For `coalesce(a, b)` where `b` is non-nullable and its evaluation emits no code (a literal or an already-evaluated value), emit a single ternary instead of the general do-while block with a global mutable `isNull` field. The empty-code requirement preserves lazy evaluation. ### Why are the changes needed? Part of SPARK-56908 (reduce generated Java size in whole-stage codegen). ~-2% of total generated source on a TPC-DS codegen dump, and one global mutable field removed per affected coalesce. ### Does this PR introduce any user-facing change? No. ### How was this patch tested? NullExpressionsSuite: new test for the ternary fast path; the SPARK-22705 test now uses three children to keep exercising the general path. Generated-by: Claude Code (Opus 4.8) --- .../expressions/nullExpressions.scala | 30 +++++++++++++++++-- .../expressions/NullExpressionsSuite.scala | 13 +++++++- 2 files changed, 40 insertions(+), 3 deletions(-) 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..578ee658c6181 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 = { + // 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(0).genCode(ctx) + val second = children(1).genCode(ctx) + if (second.code.isEmpty) { + val resultType = CodeGenerator.javaType(dataType) + 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}) { 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..cc0320e7e9e4c 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 @@ -207,10 +207,21 @@ 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("AtLeastNNonNulls should not throw 64KiB exception") { val inputs = (1 to 4000).map(x => Literal(s"x_$x")) checkEvaluation(AtLeastNNonNulls(1, inputs), true) From 3b246875151144601f9624887941ee8c5c59e9da Mon Sep 17 00:00:00 2001 From: Gengliang Wang Date: Mon, 6 Jul 2026 05:13:53 +0000 Subject: [PATCH 2/3] [SPARK-57902][SQL] Address review: add BoundReference and probe-reuse Coalesce tests Per review feedback, cover the motivating hot path directly: coalesce(nullable BoundReference, primitive literal) exercising the ternary fast path from a row read (both null and non-null), and a two-arg coalesce whose non-nullable fallback emits code, exercising the probe-reuse branch that threads the already-generated ExprCodes into the general do-while path. Generated-by: Claude Code (Opus 4.8) --- .../expressions/NullExpressionsSuite.scala | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) 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 cc0320e7e9e4c..25535246de7cf 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 @@ -20,6 +20,7 @@ package org.apache.spark.sql.catalyst.expressions import java.sql.Timestamp import org.apache.spark.{SparkFunSuite, SparkRuntimeException} +import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.FunctionIdentifier import org.apache.spark.sql.catalyst.analysis.{FunctionRegistry, SimpleAnalyzer, UnresolvedAttribute} import org.apache.spark.sql.catalyst.expressions.codegen.CodegenContext @@ -222,6 +223,26 @@ class NullExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { 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) From 2f0da82256e423c171e136c271da2180f5d70900 Mon Sep 17 00:00:00 2001 From: Gengliang Wang Date: Mon, 6 Jul 2026 16:33:02 +0000 Subject: [PATCH 3/3] [SPARK-57902][SQL] Address review: dedup resultType, use children.head, fix import order Per @LuciferYang's review: - hoist `val resultType = CodeGenerator.javaType(dataType)` to the top of doGenCode so it is computed once instead of in both the fast-path and general-path branches; - use `children.head` instead of `children(0)`; - fix the import order of InternalRow / FunctionIdentifier in NullExpressionsSuite. Generated-by: Claude Code (Opus 4.8) --- .../spark/sql/catalyst/expressions/nullExpressions.scala | 5 ++--- .../sql/catalyst/expressions/NullExpressionsSuite.scala | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) 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 578ee658c6181..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,6 +93,7 @@ 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 @@ -101,10 +102,9 @@ case class Coalesce(children: Seq[Expression]) // 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(0).genCode(ctx) + val first = children.head.genCode(ctx) val second = children(1).genCode(ctx) if (second.code.isEmpty) { - val resultType = CodeGenerator.javaType(dataType) return ev.copy( code = code""" |${first.code} @@ -134,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 25535246de7cf..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 @@ -20,8 +20,8 @@ package org.apache.spark.sql.catalyst.expressions import java.sql.Timestamp import org.apache.spark.{SparkFunSuite, SparkRuntimeException} -import org.apache.spark.sql.catalyst.InternalRow 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