The in-memory cache stores small input batches as they arrive, so a cached repartition over many map tasks holds thousands of small batches #6950

What is the problem the feature request solves?

Part of #6834.

When the cache is built from columnar input, convertColumnarBatchToCachedBatch (ArrowCachedBatchSerializer.scala:519) hands each input batch to encodeBatches (:397), which stores it as one cached batch. Nothing merges small batches first. The row path is not affected, because it cuts batches at spark.sql.inMemoryColumnarStorage.batchSize (:664).

Most native producers already emit full batches. DataFusion 55.2's FilterExec coalesces its output to 8192 rows, so a 1% filter over a 5M-row scan caches 7 batches of 8192 rows. A shuffle read does not. The native shuffle writer coalesces a map task's rows for each partition, but the reader returns every map output block as at least one batch, so a reducer that reads from M map tasks gets at least M batches. A relation cached right after a repartition over many map tasks therefore keeps batches of about rows / (maps x partitions) rows for as long as it is cached.

Measured on main (55ba50d) with Spark 4.1 and local[1] on an Apple M3 Max. The query is spark.read.parquet(src).repartition(50, $"a").cache() over 5M rows of four columns (bigint, int, double, string). The source is one file, or 200 files read either with Spark's default packing (7 splits) or with spark.sql.files.maxPartitionBytes=1m (200 splits):

Storage memory is the cached RDD's memSize, as the Storage tab shows it. With 200 map tasks the same rows take 34% more payload and 56% more memory. Each batch carries its own IPC metadata, compression framing for every buffer and a statistics row, and small buffers compress worse: about 1 KB more payload and 0.7 KB more JVM objects per batch.

Every read pays for the batch count as well. Medians of seven interleaved runs:

The first row runs no native code. The other two ran with a debug libcomet, so their absolute times are high, but both caches hold the same rows in the same 50 partitions, so the gap comes from the batch count.

Describe the potential solution

Coalesce small input batches in encodeBatches before they are serialized, up to spark.sql.inMemoryColumnarStorage.batchSize rows, the same cut the row path makes. A batch that is already that large would still be written as it is, so the common path keeps its direct write. The cost is one copy of the small batches when the cache is built, which every later read gets back.

Coalescing in the shuffle reader instead would also help native operators that read a shuffle, but a cache keeps its batches for its whole life, so it gains the most from doing it there.

Additional context

Larger batches make batch pruning coarser, since the statistics row is per batch. At 10000 rows it would match the row path and Spark's own cache format.