[andygrove]

Which issue does this PR close?

Closes #6951.

Rationale for this change

CachedBatchIpc.serialize writes each cached batch into a ChunkedByteBufferOutputStream with 1 MiB chunks. That stream allocates a whole chunk on its first write and, when it closes, copies a part-filled last chunk into an exactly sized buffer. So every batch whose payload is under 1 MiB, which is nearly every batch, cost a zeroed 1 MiB heap array and a copy of its payload. #6951 has the measurements, and the cost is highest for the many small batches that #6950 describes.

What changes are included in this PR?

• serialize sizes the stream's chunk from the message before writing it, capped at spark.comet.exec.inMemoryCache.chunkSize. A new messageSize computes what MessageSerializer.serialize writes for the batch at the start of a stream: the continuation token and length that prefix the metadata, the metadata padded so that the body starts at a multiple of 8 bytes, and the body. A message under the cap fills its one chunk exactly, so the stream allocates it once and toChunkedByteBuffer hands it back without a copy. Larger messages are still split into chunks at the cap, which is what lets one cached batch exceed 2 GiB.
• The payload's bytes are unchanged. Only the chunks they are written into differ.
• The chunkSize config's doc now says it is the largest chunk.
• CometInMemoryCacheBenchmark gets a write-only case, "in-memory cache write", for 100, 500, 2000 and 8192 rows per batch. It times CachedBatchIpc.serialize over 200 prebuilt batches of four columns (bigint, int with 10% nulls, double, 16-byte strings) under zstd level 1, and prints the heap each batch allocated, from the thread's allocation counter. --serialize-only runs just this case, and it needs no native library.

How are these changes tested?

Two new tests in CometInMemoryCacheSuite:

• "writes a payload under the chunk size into one chunk of its size" writes one uncompressed 485,224-byte batch. It checks that the payload is one chunk, and it uses the JVM's per-thread allocation counter to check that the write allocated less than the payload plus 64 KiB. With the chunk fixed at 1 MiB, the same write allocates 1,552,472 bytes.
• "computes a payload's size the way Arrow writes it" compares messageSize with the bytes MessageSerializer.serialize writes for 0 to 6 columns of int, string and boolean vectors with nulls, uncompressed and with zstd. Each column changes the metadata's length, and so the padding before the body.

Fixing the chunk at 1 MiB again fails the first test. Making the size 8 bytes short or 8 bytes over fails both. Dropping the cap fails the three existing multi-chunk tests, which split every payload into 61-byte chunks.

CometInMemoryCacheSuite, CometInMemoryCachePruningSuite, the Kryo suite and CachedBatchRowIteratorSuite pass locally on Spark 4.1 (109 tests) and 3.5 (105 tests, 4 canceled). Scalafix and the strict-warnings build pass on 3.5.

Benchmark results

The "main" arm is this branch with the chunk fixed at the configured size, which is the write path on main. The two arms alternated in one JVM, twice per run, over two runs, so each figure averages four cases. It ran on an Apple M3 Max with JDK 17.0.10 and the Spark 4.1 profile. Other sessions loaded the machine for part of the first run. Per batch is the best time of a case divided by 200.

Rows per batch	Payload	Allocated per batch, main	Allocated per batch, this PR	Write time per batch, main	Write time per batch, this PR
100	1.2 KiB	1038 KiB	15 KiB	44 µs	38 µs
500	3.9 KiB	1041 KiB	19 KiB	95 µs	83 µs
2000	13.6 KiB	1060 KiB	37 KiB	202 µs	165 µs
8192	57.8 KiB	1102 KiB	79 KiB	720 µs	689 µs

The allocation figures were the same in every round. The write times at 8192 rows moved more between rounds than the gap between the arms, so that row shows no difference. This measures the write step alone, so it says nothing yet about the build time of a whole cache. runBuildBenchmark has not been run on both sides.

[sunchao]

Reviewed the complete four-file diff from 7d8d9b1544b734c85a93cca2cc9d5380ad70accd to cd3a8c899304dec50d63cb3b133d36363cc3b112. The PR is not a draft. No introduced P1/P2 issues found within this review. There are no existing reviews, conversation comments, inline comments, or review threads with unresolved concerns.

Summary

• Prior state and problem: Each cached payload smaller than the default 1 MiB chunk incurred a full-size heap allocation followed by a copy into an exactly sized buffer. This addresses the allocation overhead described in #6951.
• Design approach: Compute the complete Arrow record-batch message size and use min(chunkSize, messageSize) when constructing Spark’s output stream.
• Correctness: The calculation matches Arrow 18.3.0’s default eight-byte prefix, metadata alignment, and aligned body length. The focused checks passed byte equality and decoded-value comparisons against the base implementation, including empty batches, zero columns, nulls, compression, and chunk boundaries. Size arithmetic remains Long until after applying the cap.
• Compatibility analysis: Checked ChunkedByteBufferOutputStream sources for Spark 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2.0. Their implementations are identical. No expression semantics, dependencies, codec defaults, or cache-format contracts change.
• Key design decisions: Keeping the existing chunk cap preserves support for payloads exceeding one JVM array. Continuing to use MessageSerializer.serialize preserves framing and buffer ownership. The reader already derives chunk offsets from actual lengths.
• Implementation sketch: CachedBatchIpc.messageSize calculates the allocation size, serialize applies it, the configuration description now identifies a maximum, and two tests check allocation and framing through the existing test helper.
• Performance: A bounded write-stage probe reduced heap allocation for a 485,224-byte payload from 1,548,008 to 500,840 bytes. Computing the size adds another metadata serialization, but the avoided allocation and copy are demonstrated. End-to-end timing improvements were not independently benchmarked.
• Design: The change is small and localized to cache serialization. Reusing Spark’s chunk stream and Arrow’s serializer keeps the mechanism straightforward without introducing another writer implementation.
• Abstraction & complexity: The sizing helper and test adapter are appropriately scoped. The helper relies on Arrow’s metadata serializer rather than duplicating FlatBuffer layout rules. No abstraction problem meeting the P1/P2 bar was identified.
• Behavioral changes worth calling out: Compared with branch-1.1 at 8ac6fccde53d96fec22184be287747eeae481cac, this intentionally removes the oversized temporary allocation and tail copy for small payloads. Serialized bytes and final chunk contents remain unchanged. Larger payloads retain capped chunks.
• Suggested improvements: No additional change meeting the P1/P2 reporting bar was identified.

Routed skills: review-comet-pr and review-comet-memory-pr, with repository AGENTS.md and the applicable contributor documentation.

Validation: The default Spark 4.1.3/JDK 17 reactor test-compile passed. A disposable focused suite using the compiled head passed all five tests: both new PR test bodies, 756 base-versus-head serialization/readback cases across none, zstd, and lz4, larger-than-1-MiB payloads, and a logical greater-than-2-GiB sizing check using only 1 MiB of Arrow storage. The working tree remains clean.

Exact-head CI at 2026-10-14 14:32 UTC: 23 checks passed, including the Spark-profile lint matrix and strict Scala warnings. Native-library build and Rust tests were still running. Another 55 checks were skipped, including Spark SQL and Iceberg jobs. No failures were reported. CI run.

Validation limits: The full native-backed cache, pruning, Kryo, and Spark SQL suites were not run locally. Runtime validation used Spark 4.1.3 only. The greater-than-2-GiB check validated size arithmetic and capping, not an actual multi-gigabyte write. CI was not yet complete.

[comphead]

Thanks @andygrove. I traced messageSize against Arrow 18.3.0's MessageSerializer.serialize and it matches for a stream starting at position 0, and since the size is only a hint to the chunk stream, a wrong value costs a copy, not correctness.

This branch now conflicts with main in CachedBatchIpc.serialize: #7225 moved the ChunkedByteBufferOutputStream block inside Utils.withRebasedOffsets, so the sized stream needs re-applying on the new nesting.

[comphead]

The numbers in the description come from a probe that is not in the PR. Could you run CometInMemoryCacheBenchmark.runBuildBenchmark before and after, or add the probe there as a write-only case, so the gain can be reproduced? I have not run it, so I do not know whether build time moves or only the allocation rate.

[andygrove]

Added the probe to CometInMemoryCacheBenchmark as a write-only case ("in-memory cache write", --serialize-only runs just it), and merged main to clear the conflict with #7225: the sized stream now sits inside Utils.withRebasedOffsets. The description has the numbers from the case rather than from the old probe.

Allocation is where it moves: about 1 MiB per batch falls to 15 KiB at 100 rows and 79 KiB at 8192. Write time per batch is 13 to 18% lower at 100 to 2000 rows and inside the noise at 8192. I have not run runBuildBenchmark on both sides, so I can't say yet what a whole cache build gains. Would you like that run before this merges?