Skip to content

[SPARK-60707][SQL] Avoid zstd writing past the output buffer when compressing Arrow batches - #59922

Closed
andygrove wants to merge 4 commits into
apache:masterfrom
andygrove:SPARK-60707
Closed

andygrove wants to merge 4 commits into
apache:masterfrom
andygrove:SPARK-60707

Conversation

@andygrove

@andygrove andygrove commented Oct 29, 2026 •

Copy link
Copy Markdown
Member

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)

…pressing Arrow batches

### What changes were proposed in this pull request?

Add `BoundedZstdCompressionCodec`, a subclass of Arrow's `ZstdCompressionCodec` that overrides only `doCompress` so zstd is given `compressBound(len)` as its destination capacity, and use it from `ArrowCompressionUtils.createCompressionCodec`. `ArrowCachedBatchSerializer.createCompressionCodec` now delegates to that method instead of keeping its own copy of the codec factory.

### 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 the buffer, but still passes `compressBound(len) + 8` as the destination capacity. On a block that does not compress, zstd can fill the whole capacity, so it can write up to 8 bytes past the end of the buffer. Fixed in Arrow by apache/arrow-java#1119 (20.0.0), which Spark cannot take on the 4.x branches.

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?

New tests in `ArrowCompressionUtilsSuite`: a sentinel-guard test that fails against Arrow's stock codec and passes with the new one, and a round-trip test.

### Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Claude Sonnet 5.5)
@andygrove

Copy link
Copy Markdown
Member Author

cc @peter-toth

@peter-toth peter-toth left a comment •

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the PR, @andygrove!

BoundedZstdCompressionCodec gives zstd compressBound(len) as its capacity behind the 8-byte length prefix, the same change apache/arrow-java#1119 makes upstream. Every Spark write path that compresses Arrow batches with zstd gets it through ArrowCompressionUtils.createCompressionCodec: ArrowConverters, the Python runners and the cached-batch serializer. The guard test fails against Arrow's stock codec as described (a 572-byte buffer for a 500-byte input) and passes here, and both suites pass. One test gap below.

Non-blocking

  • 1. Nothing pins that the factory returns the bounded codec: with createCompressionCodec switched back to new ZstdCompressionCodec(...), both new tests and ArrowCachedBatchSerializerSuite still pass, since the guard test builds the bounded codec directly. An isInstanceOf check in the round-trip test covers the path the write sides use. inline

test("SPARK-60707: zstd codec round-trips compressible and incompressible buffers") {
val allocator =
ArrowUtils.rootAllocator.newChildAllocator(this.getClass.getSimpleName, 0, Long.MaxValue)
val codec = ArrowCompressionUtils.createCompressionCodec("zstd", 3)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Finding 1. The guard test above builds RawBoundedZstdCompressionCodec directly, and this round-trip passes with either codec. I switched createCompressionCodec back to new ZstdCompressionCodec(zstdCompressionLevel), and ArrowCompressionUtilsSuite and ArrowCachedBatchSerializerSuite still passed (83 tests). With this check, the round-trip test fails on that change and passes here. ArrowCachedBatchSerializer now goes through the same factory, so it covers that path too:

Suggested change
val codec = ArrowCompressionUtils.createCompressionCodec("zstd", 3)
val codec = ArrowCompressionUtils.createCompressionCodec("zstd", 3)
// The write paths get the bounded codec, not Arrow's own.
assert(codec.isInstanceOf[BoundedZstdCompressionCodec])

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 23dc6e3: the round-trip test now asserts codec.isInstanceOf[BoundedZstdCompressionCodec]. Thanks!

@dongjoon-hyun dongjoon-hyun left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the PR, @andygrove. The fix itself looks right to me: BoundedZstdCompressionCodec.doCompress matches Arrow 19's bytecode except for the capacity, and every Spark path that writes zstd-compressed Arrow batches gets its codec from ArrowCompressionUtils.createCompressionCodec. Replaying the guard test's inputs with zstd-jni alone, this PR's capacity never wrote past the buffer for 1000 seeds, while the stock capacity did for 218 of them. I left 6 inline comments; none of them changes the fix.

Item 1 matters most: the regression test catches the original bug for one size with one seed.

Summary:

  1. [Medium] The guard test catches the original bug only at n = 500, and only with this seed (ArrowCompressionUtilsSuite.scala:116)

    782 of the seeds 0 to 999 never overrun with the stock codec, and the comment on line 115 gives the wrong cause.

  2. [Low] The incompressible inputs never reach zstd decompression (ArrowCompressionUtilsSuite.scala:148)

  3. [Low] The scaladoc of createCompressionCodec explains only the level, not the bounded codec (ArrowCompressionUtils.scala:45)

  4. [Nit] ArrowCachedBatchSerializer.createCompressionCodec is now a pass-through (ArrowCachedBatchSerializer.scala:269)

  5. [Nit] A failure inside the round-trip loop is reported as a memory leak (ArrowCompressionUtilsSuite.scala:164)

  6. [Nit] The bug is older than Arrow 18.3 (ArrowCompressionUtils.scala:55)

One note on the PR description: "with Native Memory Tracking on, the JVM aborts when the buffer is freed" holds when every buffer is its own malloc, as with arrow-memory-unsafe. Spark ships arrow-memory-netty, where buffers smaller than a chunk are slots inside a pooled chunk, so an overrun usually lands silently in the neighboring slot and NMT does not see it.

I agree with the test gap in this comment, so I did not repeat it inline.

Comment on lines +113 to +116
val random = new Random(0)
try {
// Random bytes do not compress, which is what makes zstd fill the whole destination.
Seq(1, 7, 100, 280, 500, 4096, 65536, 131072, 262144).foreach { n =>

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

1. [Medium] The guard test catches the original bug only at n = 500, and only with this seed

The guard itself works: with Arrow's stock ZstdCompressionCodec, this test fails at n = 500 as the PR description says. But that is the only size where it fails. I replayed these inputs with zstd-jni 1.5.7-23 alone (Zstd.compressUnsafe into native memory with a 64-byte guard, no Arrow or Spark code), once with Arrow 19's capacity (compressBound(n) + 8, starting 8 bytes in) and once with this PR's (compressBound(n)):

n Stock: bytes written past the buffer This PR
1, 7, 100, 280 0 0
500 8 0
4096, 65536, 131072, 262144 0 0

The comment on line 115 does not describe why n = 500 fails. Random bytes do not make zstd fill its destination: they go out as raw blocks, and for every other size zstd stops 51 to 1005 bytes before the end of the buffer. zstd reaches the end only when it tries to Huffman-code the literals and then gives up, and it tries that only when the most common byte value occurs more than (n >> 7) + 4 times. The 500 bytes drawn here happen to contain one value 8 times, against a limit of 7.

So the test depends on the seed and on the order of the sizes. With the same size list, 782 of the seeds 0 to 999 never write past the buffer with the stock codec. With seed 0, reversing the list or adding a 1000-byte case before 500 also hides the overrun. A later change to either would leave this test passing against the old doCompress.

Could the test try many seeds at n = 500, or assert as a control that the stock codec does reach the guard for the chosen input? With uniform bytes, about 20% of seeds overrun at n = 500; with about 2% of the bytes forced to 0, about 92% do. The stock overrun is at most 8 bytes, so it stays inside the 64-byte guard. Such a control would also start failing once Arrow 20.0.0 fixes the codec, which is when this class can be removed.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch, thanks for the replay. The test no longer depends on one seed at n = 500. The inputs are now random bytes with about 2% zeroed (so zstd usually tries to Huffman-code the literals and gives up), the comment gives that cause, and 100 such 500-byte inputs are all checked against BoundedZstdCompressionCodec. I also added the control you suggested: it asserts that Arrow's own ZstdCompressionCodec overruns the guard for at least one of those inputs, and notes that a failure there means Arrow fixed the codec and the class can be removed. Fixed in 23dc6e3.

try {
Seq(1, 100, 4096, 100000).foreach { n =>
val compressible = Array.tabulate[Byte](n)(i => (i % 7).toByte)
val incompressible = random.nextBytes(n)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

2. [Low] The incompressible inputs never reach zstd decompression

For random bytes, zstd's output is at least n + 9 bytes (frame header plus a raw block header), so AbstractCompressionCodec.compress drops it and stores the input raw behind a -1 prefix, and decompress then just slices it. The 1-byte compressible input takes the same path. These are the frame sizes for this test's inputs:

Input n zstd frame (bytes) Stored as
compressible 1 10 raw
compressible 100, 4096, 100000 23, 24, 27 zstd
incompressible 1, 100, 4096, 100000 10, 109, 4106, 100012 raw

So only three cases decode output from BoundedZstdCompressionCodec, and the incompressible half checks Arrow's raw fallback. Could the test assert the length prefix before decompressing, for example n for the compressible inputs of 100 bytes or more and -1 for the rest, so that each case shows which path it takes?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 23dc6e3: each case now asserts the length prefix before decompressing (n for compressible inputs of 100 bytes or more, -1 otherwise), so it is visible which path each one takes.

codecName match {
case "none" => NoCompressionCodec.INSTANCE
case "zstd" => new ZstdCompressionCodec(zstdCompressionLevel)
case "zstd" => new BoundedZstdCompressionCodec(zstdCompressionLevel)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

3. [Low] The scaladoc of createCompressionCodec explains only the level, not the bounded codec

With the copy in ArrowCachedBatchSerializer gone, the scaladoc above is the only place that says why the codecs are constructed directly, and it gives only the level: createCodec(codecType) drops it. Arrow 19 also has CompressionCodec.Factory.INSTANCE.createCodec(codecType, compressionLevel), which keeps the level but returns the stock codec (CommonsCompressionFactory calls new ZstdCompressionCodec(int)). Switching this line to that overload would pass the level tests and bring the overrun back. The isInstanceOf check suggested in this comment would catch that in tests, but the scaladoc would still suggest that only the level matters. Could it add a sentence such as:

   * For zstd it must also be a [[BoundedZstdCompressionCodec]], which neither factory overload
   * returns, not even `createCodec(codecType, compressionLevel)`, which does keep the level.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 23dc6e3: added the sentence about BoundedZstdCompressionCodec to the scaladoc, including that createCodec(codecType, compressionLevel) keeps the level but still returns the stock codec.

throw SparkException.internalError(
s"Unsupported Arrow compression codec: $other. Supported values: none, zstd, lz4")
}
ArrowCompressionUtils.createCompressionCodec(codecName.toLowerCase, compressionLevel)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

4. [Nit] ArrowCachedBatchSerializer.createCompressionCodec is now a pass-through

Both callers (lines 996 and 1103) pass conf.arrowCompressionCodec, which SQLConf already lowercases with Locale.ROOT and checks against none, zstd and lz4, so toLowerCase and the caselocale suppression around it do nothing. The other three write paths call ArrowCompressionUtils.createCompressionCodec directly. The delegation also turns ArrowCachedBatchSerializerSuite's "zstd compression level is honored by createCompressionCodec" into a copy of ArrowCompressionUtilsSuite's SPARK-57383 test: same data, same levels, same assertion.

Could both callers use ArrowCompressionUtils.createCompressionCodec and that test go? The two belong together: while the wrapper stays, that test is the only check that it passes the level through. CompressionCodec in the import on line 28 and VectorUnloader in the suite would then be unused.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed, and I checked that SQLConf lowercases and validates the value. Done in 23dc6e3: removed ArrowCachedBatchSerializer.createCompressionCodec, both callers use ArrowCompressionUtils.createCompressionCodec directly, and the duplicate level test is gone along with the now-unused imports.

}
}
} finally {
allocator.close()

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

5. [Nit] A failure inside the round-trip loop is reported as a memory leak

AbstractCompressionCodec.compress closes its input, and decompress its compressed buffer, only on success. If either throws, that buffer stays allocated, so allocator.close() here throws IllegalStateException ("Memory was leaked by query ...") and the zstd error is lost. For example, a doCompress that sets writerIndex(bytesWritten) without the 8-byte prefix keeps a truncated frame for the 100-byte compressible case, doDecompress fails with "Src size is incorrect", and the test reports the leak instead. Utils.tryWithSafeFinally { ... } { allocator.close() } would keep the first error.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 23dc6e3: the round-trip test uses Utils.tryWithSafeFinally { ... } { allocator.close() } so the first error is kept.

}

/**
* A [[ZstdCompressionCodec]] that never lets zstd write past its output buffer. Arrow 18.3 and

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

6. [Nit] The bug is older than Arrow 18.3

The same dstSize is already passed in the first zstd codec (ARROW-12110, Arrow 4.0.0), and in every release tag I checked up to 19.0.0. 18.3 and 19.0 are the versions Spark ships with this write path, so nothing is wrong for Spark, but "Arrow 18.3 and 19.0 pass ..." reads as if earlier releases were fine. Maybe "arrow-java up to 19.0 passes ..."? The PR description has the same wording.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 23dc6e3: the scaladoc now says "arrow-java up to 19.0". I updated the PR description to match.

- Assert that createCompressionCodec returns BoundedZstdCompressionCodec.
- Make the overrun regression test independent of the seed: use inputs that make zstd fill its destination, check many of them, and add a control that Arrow's own codec overruns on them.
- Check which path (zstd frame or raw) each round-trip case takes.
- Keep the first error if the round-trip test fails, instead of a leak report.
- Document the bounded codec requirement on createCompressionCodec.
- Remove the pass-through ArrowCachedBatchSerializer.createCompressionCodec and its duplicate test.

@dongjoon-hyun dongjoon-hyun left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the update, @andygrove. 23dc6e3 addresses all six comments from my earlier review, and the fix itself still looks right to me, so I'm approving.

I replayed the new guard test's exact inputs with zstd-jni 1.5.7-23 alone (no Arrow or Spark code), and the control is solid: for every seed from 0 to 299, at least 83 of the 100 inputs of 500 bytes overrun with Arrow 19's capacity, and this PR's capacity never overruns. Dropping toLowerCase along with ArrowCachedBatchSerializer.createCompressionCodec is also safe, since both callers pass conf.arrowCompressionCodec, which SQLConf lowercases and validates.

I left 7 inline comments on the tests and the PR description. None of them blocks this PR.

Summary:

  1. [Low] Only the 280- and 1000-byte inputs of this loop can reach the end of the buffer (ArrowCompressionUtilsSuite.scala:159)

    With Arrow 19's capacity, the other six sizes never overrun for seeds 0 to 299; with noisyBytes, the overrun shows up only between about 280 and 1000 bytes.

  2. [Low] The control can also fail after a zstd-jni upgrade, not only after an Arrow fix (ArrowCompressionUtilsSuite.scala:171)

  3. [Low] The PR description still describes the first commit (ArrowCachedBatchSerializer.scala:988)

  4. [Nit] writesPastEnd still closes its allocators in a plain finally (ArrowCompressionUtilsSuite.scala:152)

  5. [Nit] The prefix check reads the length in native byte order (ArrowCompressionUtilsSuite.scala:195)

  6. [Nit] doCompress throws a plain RuntimeException (ArrowCompressionUtils.scala:80)

  7. [Nit] The two Raw* codec classes sit far from the test that uses them (ArrowCompressionUtilsSuite.scala:48)

val bounded = new RawBoundedZstdCompressionCodec
val stock = new RawStockZstdCompressionCodec
val random = new Random(0)
Seq(1, 7, 100, 280, 1000, 4096, 65536, 262144).foreach { n =>

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

1. [Low] Only the 280- and 1000-byte inputs of this loop can reach the end of the buffer

The control below settles the seed question. I replayed this test's exact input sequence with zstd-jni 1.5.7-23 alone (java.util.Random behind scala.util.Random, no Arrow or Spark code): for every seed from 0 to 299, at least 83 of the 100 inputs of 500 bytes overrun with Arrow 19's capacity, and this PR's capacity never does.

This loop is different. These are the seeds, out of 0 to 299, whose input at each size overruns with Arrow 19's capacity:

n 1 7 100 280 1000 4096 65536 262144
Seeds that overrun 0 0 0 124 47 0 0 0

With seed 0, only n = 280 overruns. For independent noisyBytes(n) inputs, the overrun shows up only between about 280 and 1000 bytes: 93 of 300 seeds at 280, 279 at 500, 285 at 700, 42 at 1000, 1 at 1200, and none at 250 and below or at 1500 and above.

So six of the eight sizes check a capacity that zstd never gets close to, and "makes that the usual case" on line 108 holds only around 500 bytes. They do no harm as smoke checks, but a later change that broke the capacity only for large buffers would still pass here. Could the comment say which sizes can reach the end, or could the loop use sizes inside that range, such as 300, 500 and 700?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the table. Done in 1464f9c: the main loop now uses 280, 500 and 700 bytes with 100 inputs each, and the comment says only roughly 280 to 1000 bytes can reach the end of the buffer. The other sizes (1, 7, 100, 4096, 65536, 262144) stay as smoke checks and the comment says so.

Comment on lines +169 to +171
// Control: the inputs above do trigger the overrun in Arrow's own codec, so the checks above
// can fail. If this starts failing, Arrow has fixed the codec and BoundedZstdCompressionCodec
// can be removed.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

2. [Low] The control can also fail after a zstd-jni upgrade, not only after an Arrow fix

The control passes because, for these inputs, zstd tries to Huffman-code the literals, gives up, and on the way writes up to the end of the capacity it was given. That depends on zstd's literal heuristics, such as the (n >> 7) + 4 gate in HUF_compress_internal, which are internals that zstd does not promise to keep. A zstd-jni upgrade that changes them can make all 100 inputs stop short of the end while Spark is still on arrow-java 19.0. The test would then fail in the zstd-jni upgrade PR, and this comment would point at the wrong fix: removing BoundedZstdCompressionCodec would bring back the overstated capacity even though these inputs no longer show it.

Could the comment name both causes? For example:

    // Control: the inputs above do trigger the overrun in Arrow's own codec, so the checks above
    // can fail. This depends on how zstd compresses literals. If it starts failing after an Arrow
    // upgrade, Arrow has fixed the codec and BoundedZstdCompressionCodec can be removed; if it
    // starts failing after a zstd-jni upgrade, pick inputs that still overrun instead.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 1464f9c: the control is now per size, and its comment names both causes (Arrow upgrade: remove the class; zstd-jni upgrade: pick inputs that still overrun), using your wording.

compressionLevel: Int) extends Iterator[ArrowCachedBatch] {

private val compressionCodec = ArrowCachedBatchSerializer.createCompressionCodec(
private val compressionCodec = ArrowCompressionUtils.createCompressionCodec(

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

3. [Low] The PR description still describes the first commit

Since 23dc6e3, three parts of the description no longer match the code, and the merge script copies the description into the commit message:

  • "ArrowCachedBatchSerializer.createCompressionCodec was a second copy of the same codec factory; it now delegates to ArrowCompressionUtils.createCompressionCodec": the method is gone, and both iterators here call ArrowCompressionUtils.createCompressionCodec directly. Its level test in ArrowCachedBatchSerializerSuite is gone with it, as a copy of the SPARK-57383 test.
  • The guard test "compresses random (incompressible) inputs of several sizes from 1 byte to 256 KB" and fails against the stock codec with "zstd wrote past the end of a 572 byte output buffer (input size 500)": the inputs are now noisyBytes, with 100 of them at 500 bytes and a control on the stock codec, and that message no longer exists. By my replay, if doCompress went back to Arrow 19's capacity, the test would now fail first at n = 280, with "zstd wrote past the end of its output buffer (input size 280)".
  • "ArrowCachedBatchSerializerSuite passed locally before ArrowCachedBatchSerializer was changed to delegate to ArrowCompressionUtils; after that change it is covered by CI": there is no delegation any more.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated the description: the removed wrapper and its duplicate test, the new guard test and controls, the round-trip assertions, and the local run (82 tests across both suites).

Comment on lines +149 to +152
} finally {
guardedAllocator.close()
inputAllocator.close()
}

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

4. [Nit] writesPastEnd still closes its allocators in a plain finally

The round-trip test now keeps the first error with Utils.tryWithSafeFinally, but this helper has the same pattern. If compressRaw throws after allocating its output and leaves it open, guardedAllocator.close() throws "Memory was leaked by query" in place of the original exception. inputAllocator.close() is then skipped, so a child of the shared ArrowUtils.rootAllocator stays open for the rest of the JVM.

Today only the zstd error branch of doCompress can fail after the allocation, and it closes the buffer, so this matters only for a regression in doCompress, for example a wrong writerIndex that throws IndexOutOfBoundsException before the buffer is returned. For example:

    Utils.tryWithSafeFinally {
      ...
    } {
      Utils.tryWithSafeFinally(guardedAllocator.close())(inputAllocator.close())
    }

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 1464f9c, using the nested Utils.tryWithSafeFinally from your suggestion.

// The buffer starts with the uncompressed length, or -1 if Arrow kept the input as is
// because zstd did not make it smaller. Check which path each case takes.
val expectZstdFrame = (expected eq compressible) && n >= 100
assert(compressed.getLong(0) === (if (expectZstdFrame) n.toLong else -1L))

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

5. [Nit] The prefix check reads the length in native byte order

ArrowBuf.getLong reads in native order (MemoryUtil.getLong, which is Unsafe.getLong). AbstractCompressionCodec.writeUncompressedLength writes the prefix little-endian on every platform: it applies Long.reverseBytes when MemoryUtil.LITTLE_ENDIAN is false. On a big-endian JVM, the three zstd-frame cases would read Long.reverseBytes(n) and fail here, although the round trip itself works. The -1 cases read the same either way. If big-endian matters, reading the prefix as little-endian keeps the check portable (with import java.nio.ByteOrder):

          val prefix = compressed.nioBuffer(0, 8).order(ByteOrder.LITTLE_ENDIAN).getLong(0)
          assert(prefix === (if (expectZstdFrame) n.toLong else -1L))

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 1464f9c: the prefix is read through nioBuffer(0, 8).order(ByteOrder.LITTLE_ENDIAN).

compressionLevel)
if (Zstd.isError(bytesWritten)) {
compressedBuffer.close()
throw new RuntimeException("Error compressing: " + Zstd.getErrorName(bytesWritten))

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

6. [Nit] doCompress throws a plain RuntimeException

This line is copied from Arrow, while the rest of this file reports internal errors with SparkException.internalError. With compressBound(srcSize) as the capacity, zstd cannot run out of room, so an error here is unexpected and fits an internal error. Keeping Arrow's line unchanged has a point too, since it keeps the class easy to compare with upstream until it is removed. Either way is fine with me.

Suggested change
throw new RuntimeException("Error compressing: " + Zstd.getErrorName(bytesWritten))
throw SparkException.internalError("Error compressing: " + Zstd.getErrorName(bytesWritten))

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 1464f9c: now SparkException.internalError, which fits the rest of the file. With compressBound as the capacity zstd cannot run out of room, so an error here is unexpected.

Comment on lines +39 to +48
// Expose the protected `doCompress` so that a test can look at the raw output buffer.
private class RawBoundedZstdCompressionCodec extends BoundedZstdCompressionCodec(3) {
def compressRaw(allocator: BufferAllocator, in: ArrowBuf): ArrowBuf = doCompress(allocator, in)
}

// Arrow's own codec, which up to arrow-java 19.0 gives zstd a destination capacity that is 8
// bytes too large.
private class RawStockZstdCompressionCodec extends ZstdCompressionCodec(3) {
def compressRaw(allocator: BufferAllocator, in: ArrowBuf): ArrowBuf = doCompress(allocator, in)
}

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

7. [Nit] The two Raw* codec classes sit far from the test that uses them

Only the SPARK-60707 guard test uses RawBoundedZstdCompressionCodec and RawStockZstdCompressionCodec. Elsewhere in this suite, a helper sits right above the test that uses it: compressedSize above the SPARK-57383 test, and noisyBytes and writesPastEnd above the guard test. Moving these two classes next to writesPastEnd would keep everything the guard test needs in one place, which is also what gets deleted together on Arrow 20.0.0.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 1464f9c: both Raw* classes now sit with noisyBytes and writesPastEnd above the guard test.

- Check the sizes at which zstd can reach the end of the buffer with many inputs each, with a control per size, and name both causes of a control failure.
- Keep the first error when closing the allocators of the guard test helper.
- Read the length prefix little-endian.
- Report a zstd compression error as an internal error.
- Move the Raw codec classes next to the guard test helpers.
@dongjoon-hyun

Copy link
Copy Markdown
Member

Thank you for updating, @andygrove .

@dongjoon-hyun dongjoon-hyun left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

1464f9c addresses all seven comments from my last review, and the fix itself still looks right to me, so I'm approving. None of the comments below blocks this PR. Comments 1 and 3 are about the PR description, which the merge script copies into the commit message, so they are best settled before merging. CI has not finished on this commit yet.

I also checked the description's "compressed output is the same" with zstd-jni alone. For 29,808 inputs (four data shapes, levels -5 to 19, 50 bytes to 512 KB), zstd's output is byte-identical with compressBound(n) and compressBound(n) + 8 as the capacity, on both zstd-jni 1.5.7-23 and 1.5.7-6 (branch-4.1).

Summary:

  1. [Medium] The PR description leaves out the read failure this overrun causes (ArrowCompressionUtils.scala:60)

    With a reused root and BIGINT columns of 55 rows that do not compress, Arrow 19's capacity damaged a length prefix in 1258 to 2512 of 3000 batches (2 to 64 columns), and the reader then fails. This PR's capacity damaged none.

  2. [Low] No test goes through VectorUnloader (ArrowCompressionUtilsSuite.scala:181)

  3. [Low] "Spark cannot take a major Arrow bump on the 4.x branches" is not the direct reason (ArrowCompressionUtils.scala:61)

  4. [Nit] Correcting my earlier comment: this branch is also reached when zstd cannot allocate memory (ArrowCompressionUtils.scala:80)

  5. [Nit] expected eq compressible recomputes what each input already implies (ArrowCompressionUtilsSuite.scala:199)

/**
* A [[ZstdCompressionCodec]] that never lets zstd write past its output buffer. arrow-java up to
* 19.0 passes `compressBound(len) + 8` as the destination capacity while pointing zstd 8 bytes
* into a buffer of that size, so zstd can overrun it by up to 8 bytes; here the capacity is

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

1. [Medium] The PR description leaves out the read failure this overrun causes

The description says the overrun "lands in whatever follows the buffer, with no error". That is true while writing, but the reader can then fail. I reproduced this with Arrow 19.0.0, arrow-memory-netty and zstd-jni alone (no Spark or PR code; the right column below uses a copy of this PR's capacity): one reused VectorSchemaRoot and new VectorUnloader(root, true, codec, true), as in ArrowConverters, closing each batch before the next. The columns are BIGINT with 55 rows (440 bytes, so 8 + compressBound(440) is exactly 512), filled like noisyBytes, at level 3. These are the batches out of 3000 that had a damaged length prefix:

Columns Arrow 19's capacity This PR's capacity
1 0 0
2 1258 0
3 2508 0
4 2503 0
8 2512 0
16 2511 0
64 2505 0

The overrun from one column's output overwrites the -1 prefix of another column's raw buffer in the same batch, mostly with small values such as 1, 2 or 3, and writing reports nothing. I then passed 400 of the damaged buffers to ZstdCompressionCodec.decompress: 396 failed with Error decompressing: Unknown frame descriptor and 4 with the requested size must be non-negative.

Some limits of this result:

  • I saw no damage with either capacity for 118 rows (944 bytes) or with a new root for every batch.
  • Values changed silently only when I forced the layout by freeing the slot directly in front of a live buffer (97 of 200 tries, no error). I did not see that in the two flows above.
  • In my samples from 50 bytes to 512 KB, Arrow 19's capacity overran only at levels 1, 3 and 9, not at -5, -1 and 19.

Could the description mention the read failure? For example, right after that sentence:

With `arrow-memory-netty` this can be 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.

Buffers of 943 to 950 bytes land in that range too, but I saw no damage at 944.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the reproduction. Added to the description, right after that sentence, using your wording (including the 433 to 440 byte range and the BIGINT example).

}
}

test("SPARK-60707: zstd codec round-trips compressible and incompressible buffers") {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

2. [Low] No test goes through VectorUnloader

The guard test calls doCompress directly and never decodes its output, and the sizes in this test (1, 100, 4096, 100000) are all outside the window of about 280 to 1000 bytes, so no test runs the path where the overrun does its damage. With three BIGINT columns of 55 rows, noisyBytes-like data and a reused root, checking each data buffer's prefix after getRecordBatch() fails for 2508 of 3000 batches with Arrow 19's capacity and for none with this PR's:

val prefix = buf.nioBuffer(0, 8).order(ByteOrder.LITTLE_ENDIAN).getLong(0)
assert(prefix === -1L || prefix === rows * 8L)

It cannot replace the guard test, because how often it catches the bug depends on Netty's pool layout. It also cannot flake with BoundedZstdCompressionCodec, and it shows what users would actually hit. This one is optional.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added in 3ab1393: it writes 300 batches of three BIGINT columns with 55 rows each through ArrowWriter and VectorUnloader on a reused root, as ArrowConverters does, and checks each data buffer's prefix. With Arrow's stock ZstdCompressionCodec in the factory it fails (181 did not equal -1, and 181 did not equal 440 ... the length prefix of column 0 was overwritten), and it passes with BoundedZstdCompressionCodec. Its comment says it depends on the allocator's layout and does not replace the guard test.

* A [[ZstdCompressionCodec]] that never lets zstd write past its output buffer. arrow-java up to
* 19.0 passes `compressBound(len) + 8` as the destination capacity while pointing zstd 8 bytes
* into a buffer of that size, so zstd can overrun it by up to 8 bytes; here the capacity is
* `compressBound(len)`. Remove once on Arrow 20.0.0 (apache/arrow-java#1119).

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

3. [Low] "Spark cannot take a major Arrow bump on the 4.x branches" is not the direct reason

This is about the PR description. Arrow's major version did change within 4.x: branch-4.1 is on 18.3.0, and branch-4.2, branch-4.3 and branch-4.x are on 19.0.0. The more direct reason is that 20.0.0 is not out yet: the latest arrow-java release is v19.0.0 (2026-03-16), and on 2026-10-29 apache/arrow-java#1116 said that 20.0.0 is in preparation and that there is no capacity for patch releases. Since the description ends up in the commit message, maybe:

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.

"Remove once on Arrow 20.0.0" here is accurate as it is.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reworded in the description as suggested.

compressionLevel)
if (Zstd.isError(bytesWritten)) {
compressedBuffer.close()
throw SparkException.internalError("Error compressing: " + Zstd.getErrorName(bytesWritten))

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

4. [Nit] Correcting my earlier comment: this branch is also reached when zstd cannot allocate memory

I suggested internalError here because zstd cannot run out of room with compressBound(srcSize) as the capacity. That part holds, but I overlooked allocation: zstd-jni's compressUnsafe creates a new context and calls ZSTD_compress2, which returns memory_allocation when it cannot allocate its workspace. I read this in the zstd 1.5.7 and zstd-jni sources and did not trigger it. A native out-of-memory would then surface as INTERNAL_ERROR with Allocation error : not enough memory, which reads like a Spark bug.

The error name stays in the message, so it can still be diagnosed, and I'm fine with keeping this line as it is. Your call.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for checking. I'll keep internalError; the error name stays in the message, so a native allocation failure can still be diagnosed.

val compressed = codec.compress(allocator, input)
// The buffer starts with the uncompressed length, or -1 if Arrow kept the input as is
// because zstd did not make it smaller. Check which path each case takes.
val expectZstdFrame = (expected eq compressible) && n >= 100

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

5. [Nit] expected eq compressible recomputes what each input already implies

Nothing is wrong here. Pairing each input with its expected path removes the reference check and this val, and shows both cases on one line:

Seq(compressible -> (n >= 100), incompressible -> false).foreach {
  case (expected, expectZstdFrame) =>

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 3ab1393, using your pairing.

@dongjoon-hyun

Copy link
Copy Markdown
Member

Could you address the above, @andygrove ?

Write batches the way ArrowConverters does and check the length prefix of each data buffer, which an overrun of the previous buffer can overwrite. Also pair each round-trip input with the path it takes.
@dongjoon-hyun

Copy link
Copy Markdown
Member

Thank you so much, @andygrove !

peter-toth pushed a commit that referenced this pull request Oct 31, 2026
…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>
@peter-toth

Copy link
Copy Markdown
Contributor

Merge Summary:

Posted by merge_spark_pr.py

@peter-toth

Copy link
Copy Markdown
Contributor

Thanks @andygrove for the fix and @dongjoon-hyun for the review.

@andygrove please open backport PRs to 4.3, 4.2 and 4.1 as those were not clean cherry-picks.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants