[SPARK-59599][SQL] Pre-size the ArrowConverters.serializeBatch output buffer to avoid repeated reallocation - #58876
Open
david-mollitor-db wants to merge 2 commits into
Conversation
… buffer to avoid repeated reallocation `ArrowConverters.serializeBatch` serialized an `ArrowRecordBatch` into a `byte[]` through a `ByteArrayOutputStream` created with no initial capacity, so it started at the JDK default of 32 bytes and grew by repeatedly doubling its backing array as the batch was written. The batch is written incrementally: `MessageSerializer.serialize` writes the message metadata and then each buffer (with alignment padding) individually, and the `WritableByteChannel` returned by `Channels.newChannel` splits every write into chunks of at most 8 KB. So a 32-byte buffer reallocated and copied its contents on the order of log2(batchSize / 32) times per batch. Pre-size the buffer to `ArrowRecordBatch.computeBodyLength()` plus a small margin for the IPC message header, clamped to the `Int` range. `computeBodyLength()` is O(number of buffers) integer arithmetic over the buffer layout (no payload scan, no allocation), negligible next to serializing the body. The serialized output is byte-for-byte identical. Verified with the existing `ArrowConvertersSuite`. Co-authored-by: Isaac <no-reply@databricks.com>
…lumnar cache serializer `ArrowCachedBatchSerializer.serializeBatch` was a verbatim copy of `ArrowConverters.serializeBatch`. Widen the latter to `private[sql]` and have the columnar cache serializer delegate to it, so the Arrow columnar cache path gets the same buffer pre-sizing and the duplicated implementation (and its now-unused imports) is removed. Co-authored-by: Isaac <no-reply@databricks.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changes were proposed in this pull request?
ArrowConverters.serializeBatchserializes a singleArrowRecordBatchinto abyte[]through aByteArrayOutputStreamcreated with no initial capacity, so it starts at the JDK default of 32bytes and grows by repeatedly doubling its backing array (
grow+Arrays.copyOf) as the batchis written.
This pre-sizes the
ByteArrayOutputStreamto the batch's body length(
ArrowRecordBatch.computeBodyLength(), plus a small margin for the IPC message header), clampedto the
Intrange:ArrowCachedBatchSerializer(the Arrow columnar cache serializer) contained a verbatim copy of thesame method. This PR widens
ArrowConverters.serializeBatchtoprivate[sql]and hasArrowCachedBatchSerializer.serializeBatchdelegate to it, so the columnar cache path gets thesame pre-sizing and the duplicated implementation is removed.
Why are the changes needed?
The batch is written to the stream incrementally, not as one block:
MessageSerializer.serializewrites the message metadata and then each of the batch's buffers (with alignment padding)
individually, and the
WritableByteChannelreturned byChannels.newChannelfurther splits eachwrite into chunks of at most 8 KB. So a fresh 32-byte
ByteArrayOutputStreamreallocates andcopies its backing array on the order of
log2(batchSize / 32)times per batch as it grows,producing that many short-lived intermediate arrays (GC pressure) plus the copies.
computeBodyLength()returns the exact serialized body size and is cheap: it is O(number ofbuffers) integer arithmetic over the buffer layout (no scan of the payload, no allocation), which
is negligible next to serializing the batch body. Pre-sizing lets the buffer be allocated once at
the right size, eliminating the repeated grow-and-copy -- now on both the Arrow output path and the
Arrow columnar cache path.
This is an allocation micro-optimization; it is not expected to change throughput on its own.
Does this PR introduce any user-facing change?
No. The serialized output is byte-for-byte identical.
How was this patch tested?
Existing
ArrowConvertersSuiteandArrowCachedBatchSerializerSuitepass. This is abehavior-preserving change (buffer sizing plus de-duplication of an identical method), so no new
tests were added.
Was this patch authored or co-authored using generative AI tooling?
Generated-by: Isaac
This pull request and its description were written by Isaac.