Skip to content

Commit 073c8d8

Browse files
viiryauros-b
authored andcommitted
[SPARK-58390][SQL] Emit row counts without deserializing Arrow payloads for empty-projection cache reads
### What changes were proposed in this pull request? Short-circuit `ArrowCachedBatchSerializer.convertCachedBatchToInternalRow` when the projection is empty: emit `numRows` reused 0-field `UnsafeRow`s per cached batch, without deserializing or decompressing the batch's Arrow payload. The row count is already recorded on `ArrowCachedBatch`. Also fixes the `ArrowCachedBatch` scaladoc, which listed the per-column statistics as `(upperBound, lowerBound, ...)` while both write paths produce the `ColumnStats.collectedStatistics` order `(lowerBound, upperBound, nullCount, count, sizeInBytes)`; the code was consistently lower-first everywhere, only the doc was wrong. ### Why are the changes needed? An empty projection (e.g. a count aggregate through the row-based reader, `spark.sql.inMemoryColumnarStorage.enableVectorizedReader=false`) selects no columns, yet the reader still paid full IPC deserialization and decompression for every cached batch just to iterate its rows. That cost is pure waste: the answer is a stored integer. ### Does this PR introduce _any_ user-facing change? No. The Arrow cache serializer (SPARK-57268) is unreleased, and the change is performance-only; results are identical. ### How was this patch tested? Two new tests in `ArrowCachedBatchSerializerSuite`: - `empty projection emits row counts without deserializing the Arrow payload`: hands the reader a cached batch whose Arrow payload is garbage bytes. The empty projection returns the correct number of empty rows purely from `numRows` (fails before this change with `IllegalArgumentException: capacity < 0` from the IPC reader, proving the payload used to be deserialized), while a projection that actually needs the payload still fails on the same batch, pinning that only the empty-projection case skips the read. - `count aggregate over the cached relation with the row-based reader`: end-to-end `count(*)` over a cached relation spanning many small Arrow batches with the vectorized reader disabled, plus a `sum` over the same cached data verifying projecting reads still decode the payload correctly. Full `ArrowCachedBatchSerializerSuite` (73 tests) and `ArrowCachedBatchKryoRegistrationSuite` pass. ### Was this patch authored or co-authored using generative AI tooling? Yes, this pull request and its description were written by Claude Code. Closes #57583 from viirya/arrow-cache-empty-projection. Authored-by: Liang-Chi Hsieh <viirya@gmail.com> Signed-off-by: Uros Bojanic <221401595+uros-b@users.noreply.github.com>
1 parent b2918df commit 073c8d8

3 files changed

Lines changed: 64 additions & 4 deletions

File tree

sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatch.scala

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ import org.apache.spark.sql.columnar.SimpleMetricsCachedBatch
3333
* The batch contains:
3434
* - `numRows`: Number of rows in this batch
3535
* - `arrowData`: One encapsulated Arrow RecordBatch message (with optional compression)
36-
* - `stats`: Per-column statistics for partition pruning (upperBound, lowerBound, nullCount, etc.)
36+
* - `stats`: Per-column statistics for partition pruning (lowerBound, upperBound, nullCount, etc.)
3737
*
3838
* This format enables:
3939
* - Zero-copy columnar reads when output is ColumnarBatch with ArrowColumnVector
@@ -42,8 +42,9 @@ import org.apache.spark.sql.columnar.SimpleMetricsCachedBatch
4242
*
4343
* @param numRows Number of rows in this cached batch
4444
* @param arrowData One encapsulated Arrow RecordBatch message
45-
* @param stats Per-column statistics as InternalRow (5 fields per column:
46-
* upperBound, lowerBound, nullCount, rowCount, sizeInBytes)
45+
* @param stats Per-column statistics as InternalRow (5 fields per column, in the
46+
* `ColumnStats.collectedStatistics` order:
47+
* lowerBound, upperBound, nullCount, count, sizeInBytes)
4748
*/
4849
case class ArrowCachedBatch(
4950
numRows: Int,

sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -166,6 +166,19 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer {
166166
cacheAttributes: Seq[Attribute],
167167
selectedAttributes: Seq[Attribute],
168168
conf: SQLConf): RDD[InternalRow] = {
169+
if (selectedAttributes.isEmpty) {
170+
// Empty projection (e.g. a count aggregate over the cached relation): every cached batch
171+
// already records its row count, so emit that many empty rows without touching the Arrow
172+
// payload at all -- deserializing and decompressing it would be pure waste. The emitted
173+
// row is a single reused 0-field UnsafeRow, matching the reuse contract of the regular
174+
// path.
175+
return input.mapPartitionsInternal { batchIterator =>
176+
val rowWriter = new UnsafeRowWriter(0)
177+
rowWriter.reset()
178+
val emptyRow = rowWriter.getRow
179+
batchIterator.flatMap(batch => Iterator.fill(batch.numRows)(emptyRow))
180+
}
181+
}
169182
val cacheSchema = DataTypeUtils.fromAttributes(cacheAttributes)
170183
val selectedSchema = DataTypeUtils.fromAttributes(selectedAttributes)
171184
val timeZoneId = conf.sessionLocalTimeZone

sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializerSuite.scala

Lines changed: 47 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,10 +26,11 @@ import org.apache.arrow.vector.{
2626
TimeNanoVector, TimeStampMicroTZVector, TimeStampMicroVector, TinyIntVector,
2727
VarBinaryVector, VarCharVector, VectorSchemaRoot, VectorUnloader}
2828

29-
import org.apache.spark.{SparkConf, SparkUnsupportedOperationException}
29+
import org.apache.spark.{SparkConf, SparkException, SparkUnsupportedOperationException}
3030
import org.apache.spark.sql.{QueryTest, Row}
3131
import org.apache.spark.sql.catalyst.InternalRow
3232
import org.apache.spark.sql.catalyst.expressions.{AttributeReference, GenericInternalRow}
33+
import org.apache.spark.sql.columnar.CachedBatch
3334
import org.apache.spark.sql.execution.arrow.ArrowWriter
3435
import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf}
3536
import org.apache.spark.sql.test.{ExamplePoint, ExamplePointUDT, SharedSparkSession}
@@ -2401,6 +2402,51 @@ class ArrowCachedBatchSerializerSuite extends QueryTest with SharedSparkSession
24012402
// If the produced root was not closed, this throws "Memory was leaked by query".
24022403
alloc.close()
24032404
}
2405+
2406+
test("empty projection emits row counts without deserializing the Arrow payload") {
2407+
// A count-style read selects no columns, and the row count is already recorded on the cached
2408+
// batch, so the Arrow payload must not be deserialized at all. Prove it by handing the reader
2409+
// a batch whose payload is garbage: the empty projection succeeds purely from numRows, while
2410+
// a projection that actually needs the payload fails on the same batch.
2411+
val serializer = new ArrowCachedBatchSerializer
2412+
val attrs = Seq(AttributeReference("i", IntegerType)())
2413+
val garbage = ArrowCachedBatch(
2414+
numRows = 5,
2415+
arrowData = Array[Byte](0x13, 0x37, 0x00, -1),
2416+
stats = InternalRow(null, null, 0, 5, 4L))
2417+
val input = spark.sparkContext.parallelize(Seq[CachedBatch](garbage), 1)
2418+
val conf = spark.sessionState.conf
2419+
2420+
val emptyProjection =
2421+
serializer.convertCachedBatchToInternalRow(input, attrs, Seq.empty, conf)
2422+
assert(emptyProjection.map(_.numFields).collect() === Array(0, 0, 0, 0, 0))
2423+
2424+
val fullProjection =
2425+
serializer.convertCachedBatchToInternalRow(input, attrs, attrs, conf)
2426+
intercept[SparkException] {
2427+
fullProjection.collect()
2428+
}
2429+
}
2430+
2431+
test("count aggregate over the cached relation with the row-based reader") {
2432+
// End-to-end coverage of the empty-projection read: with the vectorized reader disabled the
2433+
// scan produces rows via convertCachedBatchToInternalRow, and a count aggregate selects no
2434+
// columns. Small Arrow batches make the count span many cached batches.
2435+
withSQLConf(
2436+
SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "false",
2437+
SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "10") {
2438+
val df = (1 to 257).toList.toDF("i")
2439+
df.cache()
2440+
try {
2441+
assert(df.count() === 257)
2442+
checkAnswer(df.selectExpr("count(*)"), Seq(Row(257)))
2443+
// A projecting query over the same cached data still reads the payload correctly.
2444+
checkAnswer(df.selectExpr("sum(i)"), Seq(Row((1 to 257).sum.toLong)))
2445+
} finally {
2446+
df.unpersist()
2447+
}
2448+
}
2449+
}
24042450
}
24052451

24062452
/**

0 commit comments

Comments
 (0)