This is an automated email from the ASF dual-hosted git repository.
gengliangwang pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 0bbefb6015fa [SPARK-57905][SQL] Collapse null-guarded UnsafeRowWriter
writes into a single writeNullable call
0bbefb6015fa is described below
commit 0bbefb6015faf69813fe593a8c38a0dc96cf3bdc
Author: Gengliang Wang <[email protected]>
AuthorDate: Mon Jul 6 09:38:28 2026 -0700
[SPARK-57905][SQL] Collapse null-guarded UnsafeRowWriter writes into a
single writeNullable call
### What changes were proposed in this pull request?
`GenerateUnsafeProjection` emits a 5-line `if (isNull) { setNullAt } else {
write }` block for every
nullable column it writes -- hitting aggregate grouping keys, exchange/sort
rows, ColumnarToRow
output, and join keys.
Before -- one block per nullable column (here an `int` column):
```java
if (isNull_0) {
rowWriter.setNullAt(0);
} else {
rowWriter.write(0, value_0);
}
```
After -- a single call; the `if/else` moves into a `writeNullable` overload
on `UnsafeRowWriter`:
```java
rowWriter.writeNullable(0, value_0, isNull_0);
```
This adds a `writeNullable(ordinal, value, isNull)` overload family to
`UnsafeRowWriter` (mirroring
the existing `write` overloads so Janino resolves the same target for a
given value type; the null
branch of each overload matches the setNull code the codegen previously
emitted for that type -- a
plain `setNullAt` for most types, but a slot-reserving `write(ordinal,
null)` for
`CalendarInterval`/`TimestampNanos` and `write(ordinal, null, precision,
scale)` for decimals) and
collapses the emission to a single call. The collapse is gated to values
that are already-evaluated
variables/literals, so non-trivial value expressions (e.g. nested-struct
field reads) keep the
`if/else` form and stay lazily evaluated; nested types (struct/array/map)
also keep the old form.
Note this is a relocation: the `if/else` still exists as bytecode, but
inside a method on the
`final` `UnsafeRowWriter` (which the JIT inlines) rather than being
re-emitted inline into every
generated projection method. That is what shrinks the generated method --
the metric behind the
HotSpot 8KB JIT-compilation threshold and the 64KB method limit -- while
total process bytecode is
roughly unchanged.
### Why are the changes needed?
Part of [SPARK-56908](https://issues.apache.org/jira/browse/SPARK-56908)
(reduce generated Java size
in whole-stage codegen). On a TPC-DS codegen dump (150 queries, 1,572
whole-stage-codegen subtrees):
about **-6.5%** of total generated source and **-6.9%** of summed per-stage
max method bytecode
(147/150 queries smaller, 0 larger) -- the metric behind the HotSpot 8KB
JIT-compilation threshold
and the 64KB method limit. The `writeNullable` branch is the same branch as
before, moved into a
method on a `final` class that the JIT inlines.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
`UnsafeRowWriterSuite` -- added a test that `writeNullable` matches the
`setNullAt`/`write` split for
null and non-null inputs across primitives, strings, wide and compact
decimals, and calendar
intervals. Existing `GeneratedProjectionSuite`,
`GenerateUnsafeProjectionSuite`,
`UnsafeRowConverterSuite`, `ExpressionEncoderSuite`, and `RowEncoderSuite`
cover nested
struct/array/map/UDT round-trips through both the collapsed and the
retained if/else write paths.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8)
Closes #56974 from gengliangwang/SPARK-57905-write-nullable.
Authored-by: Gengliang Wang <[email protected]>
Signed-off-by: Gengliang Wang <[email protected]>
(cherry picked from commit 47ca1494d015ec058da6809aebb1cae78f980de0)
Signed-off-by: Gengliang Wang <[email protected]>
---
.../expressions/codegen/UnsafeRowWriter.java | 117 +++++++++++++++++++++
.../codegen/GenerateUnsafeProjection.scala | 57 ++++++++--
.../expressions/codegen/UnsafeRowWriterSuite.scala | 63 +++++++++++
3 files changed, 229 insertions(+), 8 deletions(-)
diff --git
a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/codegen/UnsafeRowWriter.java
b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/codegen/UnsafeRowWriter.java
index 582882374e18..b6cc6425f267 100644
---
a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/codegen/UnsafeRowWriter.java
+++
b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/codegen/UnsafeRowWriter.java
@@ -21,6 +21,11 @@ import org.apache.spark.sql.catalyst.expressions.UnsafeRow;
import org.apache.spark.sql.types.Decimal;
import org.apache.spark.unsafe.Platform;
import org.apache.spark.unsafe.bitset.BitSetMethods;
+import org.apache.spark.unsafe.types.BinaryView;
+import org.apache.spark.unsafe.types.CalendarInterval;
+import org.apache.spark.unsafe.types.TimestampNanosVal;
+import org.apache.spark.unsafe.types.UTF8String;
+import org.apache.spark.unsafe.types.VariantVal;
/**
* A helper class to write data into global row buffer using `UnsafeRow`
format.
@@ -218,4 +223,116 @@ public final class UnsafeRowWriter extends UnsafeWriter {
increaseCursor(16);
}
}
+
+ // The `writeNullable` family collapses the common codegen pattern
+ // `if (isNull) { setNullAt(i); } else { write(i, value); }` into a single
call, which keeps the
+ // generated code small. The overload set mirrors the `write` overloads so
that Janino resolves
+ // the same target for a given value type. The null branch of each overload
matches the setNull
+ // code that `GenerateUnsafeProjection` emits for the corresponding data
type.
+
+ public void writeNullable(int ordinal, boolean value, boolean isNull) {
+ if (isNull) {
+ setNullAt(ordinal);
+ } else {
+ write(ordinal, value);
+ }
+ }
+
+ public void writeNullable(int ordinal, byte value, boolean isNull) {
+ if (isNull) {
+ setNullAt(ordinal);
+ } else {
+ write(ordinal, value);
+ }
+ }
+
+ public void writeNullable(int ordinal, short value, boolean isNull) {
+ if (isNull) {
+ setNullAt(ordinal);
+ } else {
+ write(ordinal, value);
+ }
+ }
+
+ public void writeNullable(int ordinal, int value, boolean isNull) {
+ if (isNull) {
+ setNullAt(ordinal);
+ } else {
+ write(ordinal, value);
+ }
+ }
+
+ public void writeNullable(int ordinal, long value, boolean isNull) {
+ if (isNull) {
+ setNullAt(ordinal);
+ } else {
+ write(ordinal, value);
+ }
+ }
+
+ public void writeNullable(int ordinal, float value, boolean isNull) {
+ if (isNull) {
+ setNullAt(ordinal);
+ } else {
+ write(ordinal, value);
+ }
+ }
+
+ public void writeNullable(int ordinal, double value, boolean isNull) {
+ if (isNull) {
+ setNullAt(ordinal);
+ } else {
+ write(ordinal, value);
+ }
+ }
+
+ public void writeNullable(int ordinal, UTF8String value, boolean isNull) {
+ if (isNull) {
+ setNullAt(ordinal);
+ } else {
+ write(ordinal, value);
+ }
+ }
+
+ public void writeNullable(int ordinal, BinaryView value, boolean isNull) {
+ if (isNull) {
+ setNullAt(ordinal);
+ } else {
+ write(ordinal, value);
+ }
+ }
+
+ public void writeNullable(int ordinal, byte[] value, boolean isNull) {
+ if (isNull) {
+ setNullAt(ordinal);
+ } else {
+ write(ordinal, value);
+ }
+ }
+
+ public void writeNullable(int ordinal, VariantVal value, boolean isNull) {
+ if (isNull) {
+ setNullAt(ordinal);
+ } else {
+ write(ordinal, value);
+ }
+ }
+
+ public void writeNullable(int ordinal, CalendarInterval value, boolean
isNull) {
+ // `write(int, CalendarInterval)` accepts a null input and still reserves
the 16-byte
+ // interval slot in the variable-length region, which is what the codegen
setNull branch
+ // relies on for this type.
+ write(ordinal, isNull ? null : value);
+ }
+
+ public void writeNullable(int ordinal, TimestampNanosVal value, boolean
isNull) {
+ // `write(int, TimestampNanosVal)` accepts a null input, see the
CalendarInterval overload.
+ write(ordinal, isNull ? null : value);
+ }
+
+ public void writeNullable(int ordinal, Decimal value, boolean isNull, int
precision, int scale) {
+ // `write(int, Decimal, int, int)` treats a null input as a null value for
both the compact
+ // and the byte-array precision ranges.
+ write(ordinal, isNull ? null : value, precision, scale);
+ }
}
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala
index 312b98093020..f729ecacb415 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala
@@ -82,6 +82,39 @@ object GenerateUnsafeProjection extends
CodeGenerator[Seq[Expression], UnsafePro
""".stripMargin
}
+ /**
+ * Returns a single `writeNullable` call that subsumes the
+ * `if (isNull) { setNullAt } else { write }` pattern, or None when the
pattern must be kept:
+ * either the data type is stored as a nested row (struct/array/map), or the
value is a
+ * non-trivial expression that must stay inside the non-null branch so it is
only evaluated for
+ * non-null inputs. Collapsing the null check keeps the generated code
small: row writes of
+ * nullable columns are emitted for every key/row written by aggregates,
sorts, exchanges, etc.
+ */
+ private def nullableWriteCall(
+ input: ExprCode,
+ index: Int,
+ dt: DataType,
+ rowWriter: String): Option[String] = {
+ val valueIsCheap = input.value match {
+ case _: VariableValue | _: GlobalValue | _: LiteralValue => true
+ case _ => false
+ }
+ if (!valueIsCheap) {
+ None
+ } else {
+ // Match the data type the same way `writeElement` does, so types that
are stored as nested
+ // rows keep the if/else form.
+ dt match {
+ case _: StructType | _: ArrayType | _: MapType | NullType => None
+ case DecimalType.Fixed(precision, scale) =>
+ Some(s"$rowWriter.writeNullable($index, ${input.value},
${input.isNull}, " +
+ s"$precision, $scale);")
+ case _ =>
+ Some(s"$rowWriter.writeNullable($index, ${input.value},
${input.isNull});")
+ }
+ }
+ }
+
private def writeExpressionsToBuffer(
ctx: CodegenContext,
row: String,
@@ -133,14 +166,22 @@ object GenerateUnsafeProjection extends
CodeGenerator[Seq[Expression], UnsafePro
|${setNull.trim}
""".stripMargin
} else {
- s"""
- |${input.code}
- |if (${input.isNull}) {
- | ${setNull.trim}
- |} else {
- | ${writeField.trim}
- |}
- """.stripMargin
+ nullableWriteCall(input, index, dt, rowWriter) match {
+ case Some(writeCall) =>
+ s"""
+ |${input.code}
+ |$writeCall
+ """.stripMargin
+ case None =>
+ s"""
+ |${input.code}
+ |if (${input.isNull}) {
+ | ${setNull.trim}
+ |} else {
+ | ${writeField.trim}
+ |}
+ """.stripMargin
+ }
}
}
diff --git
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/codegen/UnsafeRowWriterSuite.scala
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/codegen/UnsafeRowWriterSuite.scala
index 5588f42d84b2..20da970fff63 100644
---
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/codegen/UnsafeRowWriterSuite.scala
+++
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/codegen/UnsafeRowWriterSuite.scala
@@ -23,6 +23,69 @@ import org.apache.spark.unsafe.types._
class UnsafeRowWriterSuite extends SparkFunSuite {
+ test("writeNullable matches the setNullAt/write split for null and non-null
inputs") {
+ // Primitive: a null write sets the null bit, a non-null write stores the
value.
+ val intWriter = new UnsafeRowWriter(2)
+ intWriter.resetRowWriter()
+ intWriter.writeNullable(0, 42, false)
+ intWriter.writeNullable(1, -1, true)
+ val intRow = intWriter.getRow
+ assert(!intRow.isNullAt(0) && intRow.getInt(0) == 42)
+ assert(intRow.isNullAt(1))
+
+ // Reference type: the staged value is ignored when isNull is set.
+ val strWriter = new UnsafeRowWriter(1)
+ strWriter.resetRowWriter()
+ strWriter.writeNullable(0, UTF8String.fromString("ignored"), true)
+ assert(strWriter.getRow.isNullAt(0))
+
+ // Wide decimal (precision > 18): a null write must still reserve the
fixed-width slot, exactly
+ // like the write(ordinal, null, precision, scale) call the codegen
setNull branch used.
+ val decWriter = new UnsafeRowWriter(2)
+ decWriter.resetRowWriter()
+ decWriter.writeNullable(0, null, true, 38, 18)
+ val dec = Decimal(123456789.123456789)
+ assert(dec.changePrecision(38, 18))
+ decWriter.writeNullable(1, dec, false, 38, 18)
+ val decRow = decWriter.getRow
+ assert(decRow.isNullAt(0))
+ assert(decRow.getDecimal(1, 38, 18) == dec)
+
+ // Compact decimal (precision <= 18): the null bit must be set through
+ // write(ordinal, null, precision, scale)'s internal null check rather
than an explicit
+ // setNullAt branch, so it lands on the same setNullAt the codegen
previously emitted.
+ val compactDecWriter = new UnsafeRowWriter(2)
+ compactDecWriter.resetRowWriter()
+ compactDecWriter.writeNullable(0, null, true, 18, 2)
+ val compactDec = Decimal(123456789.12)
+ assert(compactDec.changePrecision(18, 2))
+ compactDecWriter.writeNullable(1, compactDec, false, 18, 2)
+ val compactDecRow = compactDecWriter.getRow
+ assert(compactDecRow.isNullAt(0))
+ assert(compactDecRow.getDecimal(1, 18, 2) == compactDec)
+
+ // CalendarInterval: a null write reserves the variable-length slot via
write(ordinal, null).
+ val intervalWriter = new UnsafeRowWriter(2)
+ intervalWriter.resetRowWriter()
+ intervalWriter.writeNullable(0, null.asInstanceOf[CalendarInterval], true)
+ val interval = new CalendarInterval(1, 2, 3L)
+ intervalWriter.writeNullable(1, interval, false)
+ val intervalRow = intervalWriter.getRow
+ assert(intervalRow.isNullAt(0))
+ assert(intervalRow.getInterval(1) == interval)
+
+ // TimestampNanosVal: like CalendarInterval, a null write must reserve the
fixed 16-byte slot
+ // via write(ordinal, null), which is the reserved-slot contract the
codegen setNull relied on.
+ val tsNanosWriter = new UnsafeRowWriter(2)
+ tsNanosWriter.resetRowWriter()
+ tsNanosWriter.writeNullable(0, null.asInstanceOf[TimestampNanosVal], true)
+ val tsNanos = new TimestampNanosVal(1234567L, 89.toShort)
+ tsNanosWriter.writeNullable(1, tsNanos, false)
+ val tsNanosRow = tsNanosWriter.getRow
+ assert(tsNanosRow.isNullAt(0))
+ assert(tsNanosRow.getTimestampNTZNanos(1) == tsNanos)
+ }
+
def checkDecimalSizeInBytes(decimal: Decimal, numBytes: Int): Unit = {
assert(decimal.toJavaBigDecimal.unscaledValue().toByteArray.length ==
numBytes)
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
