andygrove opened a new pull request, #6985: URL: https://github.com/apache/datafusion-comet/pull/6985
## 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. ## 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 three Kryo suites and `CachedBatchRowIteratorSuite` pass locally on Spark 3.4, 3.5, 4.0, 4.1 and 4.2. Scalafix and the strict-warnings build pass on 3.5. I also measured the write with a probe that isn't part of this PR. It writes 200 batches of four columns (bigint, int with 10% nulls, double and 16-byte strings) through `serialize` with zstd level 1. Each figure is the median of ten rounds, averaged over two interleaved runs per arm, on an Apple M3 Max with JDK 17 and the Spark 4.1 profile. The "main" arm is this branch with the chunk fixed at 1 MiB, which is the code path on `main`. | Rows per batch | Payload | Allocated per batch, main | Allocated per batch, this PR | Write time, main | Write time, this PR | | -------------: | ------: | ------------------------: | ---------------------------: | ---------------: | ------------------: | | 100 | 1.4 KiB | 1038 KiB | 15 KiB | 65 µs | 43 µs | | 500 | 4.4 KiB | 1041 KiB | 18 KiB | 141 µs | 92 µs | | 2000 | 15.4 KiB | 1055 KiB | 32 KiB | 210 µs | 188 µs | | 8192 | 66.5 KiB | 1109 KiB | 87 KiB | 584 µs | 544 µs | On `main` the probe ran 35 to 40 young collections for each 2000 batches it wrote. With this PR it ran none. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
