This is an automated email from the ASF dual-hosted git repository.
gengliangwang pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new c116a27e1042 [SPARK-57902][SQL] Emit a ternary for 2-arg Coalesce with
a non-nullable pure fallback
c116a27e1042 is described below
commit c116a27e104249c0a048831e2bbe0e4ca76ca050
Author: Gengliang Wang <[email protected]>
AuthorDate: Mon Jul 6 12:39:50 2026 -0700
[SPARK-57902][SQL] Emit a ternary for 2-arg Coalesce with a non-nullable
pure fallback
### What changes were proposed in this pull request?
The general `Coalesce` codegen evaluates its children inside a `do { ... }
while (false)` loop and,
before that, promotes the result's null flag to a **class-level mutable
field** via
`ctx.addMutableState`:
```scala
ev.isNull =
JavaCode.isNullGlobal(ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN,
ev.isNull))
```
It has to be a field rather than a local `boolean` because the branches may
be split out into a
separate `coalesce_N(...)` method (via
`splitExpressionsWithCurrentInputs`), so the flag must be
visible across that method boundary. So every coalesce on the general path
adds one
`private boolean ...;` field to the generated class, e.g. for `coalesce(a,
b)`:
```java
private boolean coalesce_isNull_0; // added to the generated class
...
coalesce_isNull_0 = true;
int coalesce_value_0 = -1;
do {
if (!aIsNull) { coalesce_isNull_0 = false; coalesce_value_0 = aValue;
continue; }
if (!bIsNull) { coalesce_isNull_0 = false; coalesce_value_0 = bValue;
continue; }
} while (false);
```
For `coalesce(a, b)` where `b` is non-nullable and its evaluation emits
**no** code (a literal or an
already-evaluated value) -- e.g. the `coalesce(sum, 0)` in every SUM/AVG
update -- the result can
never be null and `b` carries a pure Java expression, so the whole thing
collapses to a single
ternary with **no** field and **no** loop (`isNull` becomes the
compile-time constant
`FalseLiteral`):
```java
int coalesce_value_0 = aIsNull ? 0 : aValue;
```
The empty-code requirement preserves lazy evaluation: there are no
statements of `b` to hoist past
the null check. When the fast path does not apply, the general path reuses
the `ExprCode` already
generated while probing (generating a child twice would duplicate codegen
side effects on the
context, such as orphaned mutable state for stateful expressions and
repeated reference
registrations).
Note on why this shrinks the compiled artifact and not just the source: the
generated code is
compiled by Janino, which does no dead-code elimination or field pruning,
so the removed
`do/while` scaffold and the `private boolean` field (which also occupies a
constant-pool entry and
per-instance state) genuinely disappear from the generated class -- they
are not something a
compiler would have optimized away. Spark also intentionally bounds the
number of mutable states,
which is what the `SPARK-22705` test guards.
### Why are the changes needed?
This is part of the effort to reduce generated Java size in whole-stage
codegen (parent
[SPARK-56908](https://issues.apache.org/jira/browse/SPARK-56908)). Smaller
generated methods reduce
Janino compilation work, JIT work, and pressure against the JVM 64KB method
limit. Measured on a
TPC-DS codegen dump (150 queries, 1,572 whole-stage-codegen subtrees):
about **-2%** of total
generated source (109/150 queries smaller, 0 larger), and one global
mutable field removed per
affected coalesce.
### Does this PR introduce _any_ user-facing change?
No. This only changes the shape of generated code; results are unchanged.
### How was this patch tested?
`NullExpressionsSuite` -- added a test asserting the two-argument
non-nullable-pure-fallback form
compiles to a ternary with no global mutable state and evaluates correctly
on both the null and
non-null branches. The existing `SPARK-22705` global-variable test is
updated to use three children
so it continues to exercise the general do-while path (the two-argument
form it previously used now
takes the new fast path).
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8)
Closes #56971 from gengliangwang/SPARK-57902-coalesce-ternary.
Authored-by: Gengliang Wang <[email protected]>
Signed-off-by: Gengliang Wang <[email protected]>
---
.../sql/catalyst/expressions/nullExpressions.scala | 31 ++++++++++++++++++--
.../expressions/NullExpressionsSuite.scala | 34 +++++++++++++++++++++-
2 files changed, 61 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 e7be588c4b46..1d50d8987cce 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 5c19e69cdfa3..3d957b662d5d 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)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
