andygrove opened a new pull request, #59922:
URL: https://github.com/apache/spark/pull/59922
### What changes were proposed in this pull request?
Adds `BoundedZstdCompressionCodec`, a subclass of Arrow's
`ZstdCompressionCodec` that overrides only `doCompress` so that zstd is given
`compressBound(len)` as its destination capacity, and uses it from
`ArrowCompressionUtils.createCompressionCodec`. That method creates the
write-side codec for the Arrow IPC batches written for Python and Spark Connect
when `spark.sql.execution.arrow.compression.codec=zstd`.
`ArrowCachedBatchSerializer.createCompressionCodec` was a second copy of the
same codec factory; it now delegates to
`ArrowCompressionUtils.createCompressionCodec` so the workaround lives in one
place.
Decompression is inherited from Arrow unchanged. The class carries a note
that it can be removed once Spark is on Arrow 20.0.0.
### Why are the changes needed?
In Arrow 18.3.0 and 19.0.0, `ZstdCompressionCodec.doCompress` allocates
`compressBound(len) + 8` bytes, points zstd 8 bytes into that buffer (past the
uncompressed-length prefix), but still passes `compressBound(len) + 8` as the
destination capacity. zstd only guarantees that it stays within the capacity it
is given, and on a block that does not compress it can fill all of it, so it
can write up to 8 bytes past the end of the buffer. Arrow's allocator rounds
requests up to a power of two, which hides the overrun except when the request
lands within 8 bytes below a power of two; then the overrun lands in whatever
follows the buffer, with no error (with Native Memory Tracking on, the JVM
aborts when the buffer is freed).
This is fixed upstream in apache/arrow-java#1119 (Arrow 20.0.0). Spark
cannot take a major Arrow bump on the 4.x branches, so this works around it on
the Spark side.
Only sessions that opt in to zstd are affected; the default is `none`.
### Does this PR introduce _any_ user-facing change?
No. With zstd enabled, compressed output is the same; only the capacity
handed to zstd changes, which removes the out-of-bounds write.
### How was this patch tested?
Added to `ArrowCompressionUtilsSuite`:
- A test that allocates the codec's output from an allocator that puts a
sentinel guard directly behind each buffer, compresses random (incompressible)
inputs of several sizes from 1 byte to 256 KB, and asserts the guard is intact.
Run against Arrow's stock `ZstdCompressionCodec` it fails (`zstd wrote past the
end of a 572 byte output buffer (input size 500)`); it passes with
`BoundedZstdCompressionCodec`.
- A round-trip test for compressible and incompressible buffers through
`createCompressionCodec("zstd", ...)`.
Ran `ArrowCompressionUtilsSuite` locally on this tree.
`ArrowCachedBatchSerializerSuite` passed locally before
`ArrowCachedBatchSerializer` was changed to delegate to
`ArrowCompressionUtils`; after that change it is covered by CI.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Sonnet 5.5)
--
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]
