GH-1116: [Java] Fix compressed buffer prefix write and ZSTD dstCapacity by LuciferYang · Pull Request #1119 · apache/arrow-java

LuciferYang · 2026-04-18T13:35:04Z

What

Two fixes in the compression codec:

AbstractCompressionCodec.compress(): capture writerIndex() once

The previous code read uncompressedBuffer.writerIndex() at multiple sites — for the size comparison and again after doCompress() to populate the 8-byte uncompressed-length prefix. Capture the value once at the top of compress() and reuse it for the empty-buffer check, the size comparison, and the prefix, so all three consumers see the same value.
ZstdCompressionCodec.doCompress(): dstCapacity overstated by 8 bytes

Zstd.compressUnsafe(dst, dstSize, ...) expects dstSize to be the available space from dst. The code offsets dst by 8 bytes past the prefix but passed 8 + maxSize instead of maxSize. The compressBound() headroom hides this in practice, but the parameter was semantically wrong. Pass maxSize.

Tests

Covered by the existing round-trip tests (testEmptyBuffer, testReadWriteStream, testReadWriteFile, etc.). I was not able to construct a minimal reproducer for the original declaredUncompressed=0 symptom on the unfixed code, so both fixes are conservative correctness improvements derived from code inspection rather than failing-then-green regression tests.

lidavidm · 2026-06-29T03:45:59Z

While I think this change is fine, nothing should be modifying the buffer during this method. I don't believe any vectors/etc. are safe to use concurrently, so the comment is misleading about the guarantees we can provide.

Good point — the original framing implied a concurrency guarantee we don't actually provide, and there's no in-tree doCompress implementation that mutates the source buffer's writerIndex. Tightened the comment to state only the engineering invariant (three in-method consumers must observe the same value) and dropped the "shared reference" / doCompress speculation:

// GH-1116: capture writerIndex() once so the empty-buffer check, size // comparison, and uncompressed-length prefix all see the same value. long uncompressedLength = uncompressedBuffer.writerIndex();

Pushed in d7b6a8f.

@lidavidm

The previous comment said the captured writerIndex() guarded against the buffer being a 'shared reference to a vector's internal buffer'. As @lidavidm pointed out, Arrow vectors are not safe for concurrent use, so that framing implied a guarantee we do not provide. Tighten the comment to state only the engineering invariant: the three in-method consumers (empty-buffer check, size comparison, and the uncompressed-length prefix) must all see the same value. Reference apacheGH-1116 for context.

lidavidm · 2026-06-29T06:37:32Z

@LuciferYang were you able to replicate the original issue, by the way?

LuciferYang · 2026-06-30T02:46:50Z

Honest answer: no. I built the two new tests around the shape of data the reporter described (all-null TimestampMilli, 3469 rows, VectorSchemaRoot reuse via clear()/allocateNew()), but I just re-checked with the source change reverted — both still pass. So they're round-trip coverage for that shape, not a direct regression test for the declaredUncompressed=0 symptom.

The fix came from reading the code: the old compress() reads writerIndex() twice — once for the size comparison, once for the 8-byte prefix after doCompress(). Reading once and reusing the value is the conservative invariant. The ZSTD dstCapacity change is similar — the compressBound() headroom hides it, but the parameter was semantically wrong.

If you'd rather hold until we have a failing-then-green test, that's fair. Otherwise happy to dig further — the reporter's env was Arrow 18.3.0/19.0.0 with zstd-jni 1.5.7-6/-7, but the allocator setup wasn't in the report.

lidavidm · 2026-06-30T02:52:55Z

I think we should document that the new tests don't actually reflect the original bug report, at least (or consider if they're necessary even). But the fixes themselves are worth having.

…H-1116 The two tests added in the previous commit didn't fail on the unfixed code, so they aren't a regression test for the reported symptom. The shapes they exercise (all-null fixed-width round-trip, multi-batch ZSTD stream) are already covered by testEmptyBuffer and testReadWriteStream. Drop them per review feedback.

LuciferYang · 2026-06-30T03:11:07Z

Fair. Dropped both tests in b8e6b46 — testEmptyBuffer and testReadWriteStream already cover those round-trip shapes.

lidavidm

Thanks!

LuciferYang · 2026-06-30T03:54:33Z

Thanks for your review @lidavidm

Dan-Nedelescu · 2026-08-24T14:50:40Z

Hi @LuciferYang thanks for this fix! I'm interested in using this fix for my project. Do you know what the release schedule might be here?

LuciferYang · 2026-08-24T14:53:00Z

Sorry, I don't know. I am not a PMC Member of Apache Arrow

Dan-Nedelescu · 2026-08-24T14:55:09Z

No worries and thanks for the response!

…pressing Arrow batches ### 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 and a pass-through of the already lowercased `spark.sql.execution.arrow.compression.codec`. It is removed, and its two callers use `ArrowCompressionUtils.createCompressionCodec` directly, so the workaround lives in one place. Its level test in `ArrowCachedBatchSerializerSuite` duplicated the SPARK-57383 test in `ArrowCompressionUtilsSuite` and is removed with it. 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-java up to 19.0.0 (Spark ships 18.3.0 and 19.0.0 with this write path), `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 (when every buffer is its own `malloc`, as with `arrow-memory-unsafe`, and Native Memory Tracking is on, the JVM aborts when the buffer is freed; with the pooled chunks of `arrow-memory-netty`, which Spark ships, the overrun usually lands silently in the neighboring slot). With `arrow-memory-netty` the overrun can also hit the next buffer of the same batch: its length prefix gets overwritten, and reading the batch then fails with `Error decompressing: Unknown frame descriptor`. The request lands in that range for buffers of 433 to 440 bytes, for example 55 rows of a BIGINT column that does not compress. This is fixed upstream in apache/arrow-java#1119 (Arrow 20.0.0). Arrow 20.0.0 is not released yet, and the maintenance branches cannot take a major Arrow bump, 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 inputs with about 2% of the bytes zeroed (which makes zstd fill its whole destination for inputs of roughly 280 to 1000 bytes), and asserts the guard is intact. It checks 100 inputs each at 280, 500 and 700 bytes, and a few other sizes from 1 byte to 256 KB as smoke checks. For each of the three sizes, a control asserts that Arrow's stock `ZstdCompressionCodec` does write past the buffer for at least one input, so the check can fail. It passes with `BoundedZstdCompressionCodec`. - A test that writes batches of three BIGINT columns the way `ArrowConverters` does, with one reused root and `VectorUnloader`, and asserts that no data buffer's length prefix was overwritten by the previous buffer. With Arrow's stock `ZstdCompressionCodec` in the factory it fails; with `BoundedZstdCompressionCodec` it passes. - A round-trip test for compressible and incompressible buffers through `createCompressionCodec("zstd", ...)`. It asserts that the factory returns `BoundedZstdCompressionCodec` and which path (zstd frame or raw) each case takes. Ran `ArrowCompressionUtilsSuite` and `ArrowCachedBatchSerializerSuite` locally on this tree (83 tests, all passing). ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Claude Sonnet 5.5) Closes #59922 from andygrove/SPARK-60707. Authored-by: Andy Grove <agrove@apache.org> Signed-off-by: Peter Toth <peter.toth@gmail.com>

…pressing Arrow batches ### 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 and a pass-through of the already lowercased `spark.sql.execution.arrow.compression.codec`. It is removed, and its two callers use `ArrowCompressionUtils.createCompressionCodec` directly, so the workaround lives in one place. Its level test in `ArrowCachedBatchSerializerSuite` duplicated the SPARK-57383 test in `ArrowCompressionUtilsSuite` and is removed with it. 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-java up to 19.0.0 (Spark ships 18.3.0 and 19.0.0 with this write path), `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 (when every buffer is its own `malloc`, as with `arrow-memory-unsafe`, and Native Memory Tracking is on, the JVM aborts when the buffer is freed; with the pooled chunks of `arrow-memory-netty`, which Spark ships, the overrun usually lands silently in the neighboring slot). With `arrow-memory-netty` the overrun can also hit the next buffer of the same batch: its length prefix gets overwritten, and reading the batch then fails with `Error decompressing: Unknown frame descriptor`. The request lands in that range for buffers of 433 to 440 bytes, for example 55 rows of a BIGINT column that does not compress. This is fixed upstream in apache/arrow-java#1119 (Arrow 20.0.0). Arrow 20.0.0 is not released yet, and the maintenance branches cannot take a major Arrow bump, 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 inputs with about 2% of the bytes zeroed (which makes zstd fill its whole destination for inputs of roughly 280 to 1000 bytes), and asserts the guard is intact. It checks 100 inputs each at 280, 500 and 700 bytes, and a few other sizes from 1 byte to 256 KB as smoke checks. For each of the three sizes, a control asserts that Arrow's stock `ZstdCompressionCodec` does write past the buffer for at least one input, so the check can fail. It passes with `BoundedZstdCompressionCodec`. - A test that writes batches of three BIGINT columns the way `ArrowConverters` does, with one reused root and `VectorUnloader`, and asserts that no data buffer's length prefix was overwritten by the previous buffer. With Arrow's stock `ZstdCompressionCodec` in the factory it fails; with `BoundedZstdCompressionCodec` it passes. - A round-trip test for compressible and incompressible buffers through `createCompressionCodec("zstd", ...)`. It asserts that the factory returns `BoundedZstdCompressionCodec` and which path (zstd frame or raw) each case takes. Ran `ArrowCompressionUtilsSuite` and `ArrowCachedBatchSerializerSuite` locally on this tree (83 tests, all passing). ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Claude Sonnet 5.5) Closes #59922 from andygrove/SPARK-60707. Authored-by: Andy Grove <agrove@apache.org> Signed-off-by: Peter Toth <peter.toth@gmail.com> (cherry picked from commit 20ea29c) Signed-off-by: Peter Toth <peter.toth@gmail.com>

fix

4c754b1

LuciferYang requested review from jbonofre, laurentgo, lidavidm and wgtmac as code owners April 18, 2026 13:35

LuciferYang changed the title ~~# GH-1116: [Java] Fix compressed buffer prefix write and ZSTD dstCapacity~~ GH-1116: [Java] Fix compressed buffer prefix write and ZSTD dstCapacity Apr 18, 2026

This comment has been minimized.

Sign in to view

lidavidm added the bug-fix PRs that fix a big. label Jun 29, 2026

github-actions Bot added this to the 20.0.0 milestone Jun 29, 2026

lidavidm reviewed Jun 29, 2026

View reviewed changes

lidavidm approved these changes Jun 30, 2026

View reviewed changes

lidavidm merged commit b0e51af into apache:main Jun 30, 2026
25 of 26 checks passed

andygrove mentioned this pull request Oct 31, 2026

[EPIC] Track arrow-java bugs and limitations that Comet works around apache/datafusion-comet#7362

Open

15 tasks

Navigation Menu

Repository navigation

GH-1116: [Java] Fix compressed buffer prefix write and ZSTD dstCapacity - #1119

GH-1116: [Java] Fix compressed buffer prefix write and ZSTD dstCapacity#1119
lidavidm merged 3 commits into
apache:mainfrom
LuciferYang:fix-zstd-compressed-buffer-prefix-corruption

LuciferYang commented Apr 18, 2026 •

edited by lidavidm

Loading

Uh oh!

This comment has been minimized.

lidavidm Jun 29, 2026

Uh oh!

LuciferYang Jun 29, 2026

Uh oh!

lidavidm commented Jun 29, 2026

Uh oh!

LuciferYang commented Jun 30, 2026

Uh oh!

lidavidm commented Jun 30, 2026 •

edited

Loading

Uh oh!

LuciferYang commented Jun 30, 2026

Uh oh!

lidavidm left a comment

Uh oh!

LuciferYang commented Jun 30, 2026

Uh oh!

Uh oh!

Dan-Nedelescu commented Aug 24, 2026

Uh oh!

LuciferYang commented Aug 24, 2026

Uh oh!

Dan-Nedelescu commented Aug 24, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

LuciferYang commented Apr 18, 2026 • edited by lidavidm Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What

Tests

Uh oh!

This comment has been minimized.

lidavidm Jun 29, 2026

Choose a reason for hiding this comment

Uh oh!

LuciferYang Jun 29, 2026

Choose a reason for hiding this comment

Uh oh!

lidavidm commented Jun 29, 2026

Uh oh!

LuciferYang commented Jun 30, 2026

Uh oh!

lidavidm commented Jun 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

LuciferYang commented Jun 30, 2026

Uh oh!

lidavidm left a comment

Choose a reason for hiding this comment

Uh oh!

LuciferYang commented Jun 30, 2026

Uh oh!

Uh oh!

Dan-Nedelescu commented Aug 24, 2026

Uh oh!

LuciferYang commented Aug 24, 2026

Uh oh!

Dan-Nedelescu commented Aug 24, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

LuciferYang commented Apr 18, 2026 •

edited by lidavidm

Loading

lidavidm commented Jun 30, 2026 •

edited

Loading