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 758afe7de027 [SPARK-57909][SQL] Extract ColumnarToRow batch-advance
machinery into a compiled helper
758afe7de027 is described below
commit 758afe7de02716cdc8f174a337f23de4ef77dd65
Author: Gengliang Wang <[email protected]>
AuthorDate: Mon Jul 6 09:25:04 2026 -0700
[SPARK-57909][SQL] Extract ColumnarToRow batch-advance machinery into a
compiled helper
### What changes were proposed in this pull request?
Every `ColumnarToRowExec` stage re-emits the same type-independent
batch-advance machinery into its
generated `nextBatch` method -- release the current batch, fetch the next
one from the input
iterator, bump the `numInputBatches`/`numOutputRows` metrics -- plus
release lines in the
`processNext` loop epilogue.
Before -- the fetch/metrics logic is inline in `nextBatch`, and the release
is a separate pair of
lines in the `processNext` loop:
```java
private void nextBatch() throws java.io.IOException {
if (input.hasNext()) {
batch = (ColumnarBatch) input.next();
numInputBatches.add(1);
numOutputRows.add(batch.numRows());
idx = 0;
colInstance0 = (OnHeapColumnVector) batch.column(0); // per-column
casts (schema-specific)
...
}
}
// ... inside processNext()'s loop epilogue:
idx = numRows;
batch.closeIfFreeable();
batch = null;
nextBatch();
```
After -- the release/fetch/metrics move into a compiled
`ColumnarToRowExec.advanceBatch(input, current, numInputBatches,
numOutputRows)` helper; the
generated `nextBatch` keeps only the per-column `ColumnVector` casts (the
schema-specific part), and
the two release lines drop out of the loop epilogue since `advanceBatch`
now does the release:
```java
private void nextBatch() throws java.io.IOException {
batch = org.apache.spark.sql.execution.ColumnarToRowExec.advanceBatch(
input, batch, numInputBatches, numOutputRows);
if (batch != null) {
idx = 0;
colInstance0 = (OnHeapColumnVector) batch.column(0); // per-column
casts (schema-specific)
...
}
}
// ... loop epilogue is now just:
idx = numRows;
nextBatch();
```
This is a relocation: the batch-advance logic still exists as bytecode, but
in one precompiled
helper rather than re-emitted into every stage's `processNext`/`nextBatch`.
Semantics are preserved
point-by-point:
- same release-before-fetch order (`closeIfFreeable` on the previous batch
before fetching; no
release on the first call when the batch is null);
- exhaustion leaves the batch null, so the processing loop exits exactly as
before;
- mid-batch `shouldStop` re-entry is untouched (no advance happens on that
path);
- the limit-reached cleanup block is unchanged.
Exception path: if the input iterator throws mid-advance (after the helper
released the current
batch), the generated batch field briefly keeps referencing the released
batch. That state is
unobservable -- the cleanup block is not in a `finally` so it is skipped on
a throw, a failed task's
iterator is never pumped again, and reader-owned batches are not freeable
so `closeIfFreeable`
cannot release them in the first place. This reasoning is documented in the
emitter.
### 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 **-2.2%** summed per-stage max method bytecode (the advance code sits
in `processNext`,
typically the largest generated method, which is the metric behind the
HotSpot 8KB JIT-compilation
threshold) and **-0.9%** total generated source, with every query shrinking.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing tests: `WholeStageCodegenSuite`, `ParquetV1QuerySuite`, and
`ParquetV2QuerySuite` exercise
execution through the parquet read paths (batch release, exhaustion, and
re-entry).
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8)
Closes #56977 from gengliangwang/SPARK-57909-columnar-advance-batch.
Authored-by: Gengliang Wang <[email protected]>
Signed-off-by: Gengliang Wang <[email protected]>
(cherry picked from commit c9a910486056c3f267ec4f4f4340e7c800a25aa0)
Signed-off-by: Gengliang Wang <[email protected]>
---
.../org/apache/spark/sql/execution/Columnar.scala | 42 +++++++--
.../sql/execution/ColumnarToRowExecSuite.scala | 104 +++++++++++++++++++++
2 files changed, 140 insertions(+), 6 deletions(-)
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/Columnar.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/Columnar.scala
index 5c98e3685327..c6323575c9a8 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/Columnar.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/Columnar.scala
@@ -152,14 +152,19 @@ case class ColumnarToRowExec(child: SparkPlan)
(name, s"$name = ($columnVectorClz) $batch.column($i);")
}.unzip
+ // The type-independent advance machinery (releasing the current batch,
fetching the next one
+ // and bumping the metrics) is a compiled helper; only the per-column
vector casts are emitted.
+ // Exception-path note: if the helper throws after releasing the current
batch (the input
+ // iterator failing mid-scan), the field briefly keeps referencing the
released batch. That
+ // state is unobservable: the cleanup block below is not in a finally, a
failed task's
+ // iterator is never pumped again, and reader-owned batches are not
freeable to begin with.
val nextBatch = ctx.freshName("nextBatch")
val nextBatchFuncName = ctx.addNewFunction(nextBatch,
s"""
|private void $nextBatch() throws java.io.IOException {
- | if ($input.hasNext()) {
- | $batch = ($columnarBatchClz)$input.next();
- | $numInputBatches.add(1);
- | $numOutputRows.add($batch.numRows());
+ | $batch =
org.apache.spark.sql.execution.ColumnarToRowExec.advanceBatch(
+ | $input, $batch, $numInputBatches, $numOutputRows);
+ | if ($batch != null) {
| $idx = 0;
| ${columnAssigns.mkString("", "\n", "\n")}
| }
@@ -191,8 +196,6 @@ case class ColumnarToRowExec(child: SparkPlan)
| $shouldStop
| }
| $idx = $numRows;
- | $batch.closeIfFreeable();
- | $batch = null;
| $nextBatchFuncName();
|}
|// clean up resources
@@ -210,6 +213,33 @@ case class ColumnarToRowExec(child: SparkPlan)
copy(child = newChild)
}
+object ColumnarToRowExec {
+ /**
+ * Releases the current batch (if any) and fetches the next one from the
input iterator,
+ * bumping the batch/row metrics. Returns null when the input is exhausted.
This is called by
+ * the generated code of [[ColumnarToRowExec]] (through the static
forwarder), so the
+ * type-independent per-batch bookkeeping is compiled once per JVM instead
of being re-emitted
+ * into every stage's generated `nextBatch` method.
+ */
+ def advanceBatch(
+ input: Iterator[ColumnarBatch],
+ current: ColumnarBatch,
+ numInputBatches: SQLMetric,
+ numOutputRows: SQLMetric): ColumnarBatch = {
+ if (current != null) {
+ current.closeIfFreeable()
+ }
+ if (input.hasNext) {
+ val batch = input.next()
+ numInputBatches.add(1)
+ numOutputRows.add(batch.numRows())
+ batch
+ } else {
+ null
+ }
+ }
+}
+
/**
* Provides an optimized set of APIs to append row based data to an array of
* [[WritableColumnVector]].
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/ColumnarToRowExecSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/ColumnarToRowExecSuite.scala
new file mode 100644
index 000000000000..0987c7cb02f6
--- /dev/null
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/ColumnarToRowExecSuite.scala
@@ -0,0 +1,104 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.types.{Decimal, IntegerType}
+import org.apache.spark.sql.vectorized.{ColumnarArray, ColumnarBatch,
ColumnarMap, ColumnVector}
+import org.apache.spark.unsafe.types.UTF8String
+
+class ColumnarToRowExecSuite extends SparkFunSuite {
+
+ /** A minimal ColumnVector that records whether `closeIfFreeable`/`close`
was called. */
+ private class RecordingColumnVector extends ColumnVector(IntegerType) {
+ var closeIfFreeableCalls = 0
+ var closeCalls = 0
+
+ override def closeIfFreeable(): Unit = closeIfFreeableCalls += 1
+ override def close(): Unit = closeCalls += 1
+
+ private def unused = throw new UnsupportedOperationException("not used by
advanceBatch")
+ override def hasNull: Boolean = unused
+ override def numNulls(): Int = unused
+ override def isNullAt(rowId: Int): Boolean = unused
+ override def getBoolean(rowId: Int): Boolean = unused
+ override def getByte(rowId: Int): Byte = unused
+ override def getShort(rowId: Int): Short = unused
+ override def getInt(rowId: Int): Int = unused
+ override def getLong(rowId: Int): Long = unused
+ override def getFloat(rowId: Int): Float = unused
+ override def getDouble(rowId: Int): Double = unused
+ override def getArray(rowId: Int): ColumnarArray = unused
+ override def getMap(ordinal: Int): ColumnarMap = unused
+ override def getDecimal(rowId: Int, precision: Int, scale: Int): Decimal =
unused
+ override def getUTF8String(rowId: Int): UTF8String = unused
+ override def getBinary(rowId: Int): Array[Byte] = unused
+ override def getChild(ordinal: Int): ColumnVector = unused
+ }
+
+ private def newBatch(numRows: Int): (ColumnarBatch, RecordingColumnVector) =
{
+ val vector = new RecordingColumnVector
+ val batch = new ColumnarBatch(Array[ColumnVector](vector), numRows)
+ (batch, vector)
+ }
+
+ test("advanceBatch releases the current batch and fetches the next, bumping
metrics") {
+ val numInputBatches = new SQLMetric("sum")
+ val numOutputRows = new SQLMetric("sum")
+ val (current, currentVector) = newBatch(3)
+ val (next, _) = newBatch(5)
+
+ val result = ColumnarToRowExec.advanceBatch(
+ Iterator(next), current, numInputBatches, numOutputRows)
+
+ assert(result eq next)
+ // The previous batch is released before the next is fetched.
+ assert(currentVector.closeIfFreeableCalls == 1)
+ assert(numInputBatches.value == 1)
+ assert(numOutputRows.value == 5)
+ }
+
+ test("advanceBatch does not release when the current batch is null (first
call)") {
+ val numInputBatches = new SQLMetric("sum")
+ val numOutputRows = new SQLMetric("sum")
+ val (next, _) = newBatch(4)
+
+ val result = ColumnarToRowExec.advanceBatch(
+ Iterator(next), null, numInputBatches, numOutputRows)
+
+ assert(result eq next)
+ assert(numInputBatches.value == 1)
+ assert(numOutputRows.value == 4)
+ }
+
+ test("advanceBatch returns null and bumps no metrics when the input is
exhausted") {
+ val numInputBatches = new SQLMetric("sum")
+ val numOutputRows = new SQLMetric("sum")
+ val (current, currentVector) = newBatch(3)
+
+ val result = ColumnarToRowExec.advanceBatch(
+ Iterator.empty, current, numInputBatches, numOutputRows)
+
+ assert(result == null)
+ // The current batch is still released even when there is no next batch.
+ assert(currentVector.closeIfFreeableCalls == 1)
+ assert(numInputBatches.value == 0)
+ assert(numOutputRows.value == 0)
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
