Every cached batch allocates a 1 MiB heap chunk, however small its payload #6951

What is the problem the feature request solves?

Part of #6834.

CachedBatchIpc.serialize writes each cached batch into new ChunkedByteBufferOutputStream(chunkSize, ByteBuffer.allocate) (CachedBatchIpc.scala:237). chunkSize comes from the internal spark.comet.exec.inMemoryCache.chunkSize, 1 MiB by default (CometConf.scala:440). Spark's stream allocates a whole chunk on the first write (allocateNewChunkIfNeeded, ChunkedByteBufferOutputStream.scala:85 in 4.1.3, unchanged from 3.4.3 to 4.2.0). toChunkedByteBuffer then copies the last chunk into an exactly sized buffer unless it is full (:114). So every batch whose payload is under 1 MiB, which is nearly every batch, costs a zeroed 1 MiB heap array that is garbage as soon as the batch is written, plus a copy of its whole payload.

Measured on the write alone (CometCachedBatchHelper.serialize, zstd level 1, four columns of bigint, int, double and string), on an Apple M3 Max with JDK 17 and G1 at a 4 GiB heap, median of ten rounds of 200 batches:

With 1 MiB chunks the JVM ran a young collection about every 200 batches, and with exact chunks it ran none. The time per batch went from 44 to 57 µs at 100 rows, from 68 to 107 µs at 500 rows and from 149 to 169 µs at 2000 rows. At 8192 rows the difference was within run-to-run noise. A 4 MiB G1 region size, where a 1 MiB array is no longer a humongous allocation, gave a similar gap for the small batches.

End to end it matters most for small batches. The 200-map cache in #6950 (10000 batches of about 500 rows) allocates about 10 GiB of chunks to hold 37 MiB. Its build spent 160 ms in GC with 1 MiB chunks and 9 ms with 16 KiB chunks, and took 14.2 s against 13.8 s (median of five, with a debug libcomet that inflates the native part of the build). The same rows in 633 full batches spent 16 ms against 3 ms in GC, with no measurable difference in build time.

Describe the potential solution

Size the stream from the message before writing it. The compressed record batch's computeBodyLength(), plus its serialized metadata (MessageSerializer.serializeMetadata) and the 8-byte prefix with padding, is the payload size. Using min(chunkSize, payload size) as the chunk size gives a payload under 1 MiB one exactly sized chunk with no tail copy, and keeps 1 MiB chunks for the large batches that chunking exists for.

Additional context

Chunking is what lets a single cached batch exceed the 2 GiB one JVM array can hold, so it should stay for large payloads.