dongjoon-hyun commented on code in PR #59922:
URL: https://github.com/apache/spark/pull/59922#discussion_r4424301122
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtils.scala:
##########
@@ -34,17 +36,50 @@ private[sql] object ArrowCompressionUtils {
* carry a compression level, so that factory overload always builds a codec
at the default
* level, silently dropping the configured one. The level only matters on
the write side; the
* read side looks up the codec by the type recorded in the IPC message, so
it is unaffected.
+ *
+ * For zstd it must also be a [[BoundedZstdCompressionCodec]], which neither
factory overload
+ * returns, not even `createCodec(codecType, compressionLevel)`, which does
keep the level.
*/
def createCompressionCodec(
codecName: String,
zstdCompressionLevel: Int): CompressionCodec = {
codecName match {
case "none" => NoCompressionCodec.INSTANCE
- case "zstd" => new ZstdCompressionCodec(zstdCompressionLevel)
+ case "zstd" => new BoundedZstdCompressionCodec(zstdCompressionLevel)
case "lz4" => new Lz4CompressionCodec()
case other =>
throw SparkException.internalError(
s"Unsupported Arrow compression codec: $other. Supported values:
none, zstd, lz4")
}
}
}
+
+/**
+ * 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
Review Comment:
**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.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtils.scala:
##########
@@ -34,17 +36,50 @@ private[sql] object ArrowCompressionUtils {
* carry a compression level, so that factory overload always builds a codec
at the default
* level, silently dropping the configured one. The level only matters on
the write side; the
* read side looks up the codec by the type recorded in the IPC message, so
it is unaffected.
+ *
+ * For zstd it must also be a [[BoundedZstdCompressionCodec]], which neither
factory overload
+ * returns, not even `createCodec(codecType, compressionLevel)`, which does
keep the level.
*/
def createCompressionCodec(
codecName: String,
zstdCompressionLevel: Int): CompressionCodec = {
codecName match {
case "none" => NoCompressionCodec.INSTANCE
- case "zstd" => new ZstdCompressionCodec(zstdCompressionLevel)
+ case "zstd" => new BoundedZstdCompressionCodec(zstdCompressionLevel)
case "lz4" => new Lz4CompressionCodec()
case other =>
throw SparkException.internalError(
s"Unsupported Arrow compression codec: $other. Supported values:
none, zstd, lz4")
}
}
}
+
+/**
+ * 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).
+ */
+private[sql] class BoundedZstdCompressionCodec(compressionLevel: Int)
+ extends ZstdCompressionCodec(compressionLevel) {
+
+ override protected def doCompress(
+ allocator: BufferAllocator,
+ uncompressedBuffer: ArrowBuf): ArrowBuf = {
+ val srcSize = uncompressedBuffer.writerIndex()
+ val maxSize = Zstd.compressBound(srcSize)
+ val compressedBuffer =
allocator.buffer(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH + maxSize)
+ val bytesWritten = Zstd.compressUnsafe(
+ compressedBuffer.memoryAddress() +
CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH,
+ maxSize,
+ uncompressedBuffer.memoryAddress(),
+ srcSize,
+ compressionLevel)
+ if (Zstd.isError(bytesWritten)) {
+ compressedBuffer.close()
+ throw SparkException.internalError("Error compressing: " +
Zstd.getErrorName(bytesWritten))
Review Comment:
**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.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtilsSuite.scala:
##########
@@ -85,4 +91,126 @@ class ArrowCompressionUtilsSuite extends SparkFunSuite {
}
assert(e.getMessage.contains("Unsupported Arrow compression codec:
snappy"))
}
+
+ // 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)
+ }
+
+ // Random bytes with about 2% of them zeroed. Pure random bytes only make
zstd fill its whole
+ // destination when one byte value happens to be frequent enough (more than
about n / 128 + 4
+ // occurrences) for zstd to try Huffman-coding the literals and give up.
Zeroing some bytes
+ // makes that the usual case instead of a property of the seed.
+ private def noisyBytes(random: Random, n: Int): Array[Byte] = {
+ val bytes = random.nextBytes(n)
+ (0 until n).foreach { i => if (random.nextInt(50) == 0) bytes(i) = 0 }
+ bytes
+ }
+
+ // Runs `compressRaw` into a buffer that is followed by sentinel bytes, and
returns whether the
+ // sentinel was overwritten, i.e. whether zstd wrote past the end of the
buffer.
+ private def writesPastEnd(
+ compressRaw: (BufferAllocator, ArrowBuf) => ArrowBuf,
+ data: Array[Byte]): Boolean = {
+ val guardSize = 64
+ val guardByte = 0x5A.toByte
+ var requestedSize = 0L
+ val guardedAllocator = new RootAllocator(Long.MaxValue) {
+ override def buffer(size: Long): ArrowBuf = {
+ requestedSize = size
+ val buf = super.buffer(size + guardSize)
+ buf.setBytes(size, Array.fill(guardSize)(guardByte))
+ buf
+ }
+ }
+ val inputAllocator =
+ ArrowUtils.rootAllocator.newChildAllocator(this.getClass.getSimpleName,
0, Long.MaxValue)
+ Utils.tryWithSafeFinally {
+ val input = inputAllocator.buffer(data.length)
+ val out = try {
+ input.setBytes(0, data)
+ input.writerIndex(data.length)
+ compressRaw(guardedAllocator, input)
+ } finally {
+ input.close()
+ }
+ try {
+ val guard = new Array[Byte](guardSize)
+ out.getBytes(requestedSize, guard)
+ !guard.forall(_ == guardByte)
+ } finally {
+ out.close()
+ }
+ } {
+
Utils.tryWithSafeFinally(guardedAllocator.close())(inputAllocator.close())
+ }
+ }
+
+ test("SPARK-60707: zstd compression does not write past the end of its
output buffer") {
+ val bounded = new RawBoundedZstdCompressionCodec
+ val stock = new RawStockZstdCompressionCodec
+ val random = new Random(0)
+ // Of these inputs, only those of roughly 280 to 1000 bytes make zstd fill
its whole
+ // destination; smaller and larger ones stop short of the end. The sizes
in that range are
+ // checked with many inputs each, and the other sizes only as smoke checks.
+ Seq(280, 500, 700).foreach { n =>
+ val inputs = Seq.fill(100)(noisyBytes(random, n))
+ inputs.foreach { data =>
+ assert(!writesPastEnd(bounded.compressRaw, data),
+ s"zstd wrote past the end of its output buffer (input size $n)")
+ }
+ // Control: these inputs 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.
+ assert(inputs.exists(writesPastEnd(stock.compressRaw, _)),
+ s"no input of size $n made Arrow's ZstdCompressionCodec write past its
output buffer")
+ }
+ Seq(1, 7, 100, 4096, 65536, 262144).foreach { n =>
+ assert(!writesPastEnd(bounded.compressRaw, noisyBytes(random, n)),
+ s"zstd wrote past the end of its output buffer (input size $n)")
+ }
+ }
+
+ test("SPARK-60707: zstd codec round-trips compressible and incompressible
buffers") {
Review Comment:
**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:
```scala
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.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtils.scala:
##########
@@ -34,17 +36,50 @@ private[sql] object ArrowCompressionUtils {
* carry a compression level, so that factory overload always builds a codec
at the default
* level, silently dropping the configured one. The level only matters on
the write side; the
* read side looks up the codec by the type recorded in the IPC message, so
it is unaffected.
+ *
+ * For zstd it must also be a [[BoundedZstdCompressionCodec]], which neither
factory overload
+ * returns, not even `createCodec(codecType, compressionLevel)`, which does
keep the level.
*/
def createCompressionCodec(
codecName: String,
zstdCompressionLevel: Int): CompressionCodec = {
codecName match {
case "none" => NoCompressionCodec.INSTANCE
- case "zstd" => new ZstdCompressionCodec(zstdCompressionLevel)
+ case "zstd" => new BoundedZstdCompressionCodec(zstdCompressionLevel)
case "lz4" => new Lz4CompressionCodec()
case other =>
throw SparkException.internalError(
s"Unsupported Arrow compression codec: $other. Supported values:
none, zstd, lz4")
}
}
}
+
+/**
+ * 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).
Review Comment:
**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](https://github.com/apache/arrow-java/issues/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.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtilsSuite.scala:
##########
@@ -85,4 +91,126 @@ class ArrowCompressionUtilsSuite extends SparkFunSuite {
}
assert(e.getMessage.contains("Unsupported Arrow compression codec:
snappy"))
}
+
+ // 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)
+ }
+
+ // Random bytes with about 2% of them zeroed. Pure random bytes only make
zstd fill its whole
+ // destination when one byte value happens to be frequent enough (more than
about n / 128 + 4
+ // occurrences) for zstd to try Huffman-coding the literals and give up.
Zeroing some bytes
+ // makes that the usual case instead of a property of the seed.
+ private def noisyBytes(random: Random, n: Int): Array[Byte] = {
+ val bytes = random.nextBytes(n)
+ (0 until n).foreach { i => if (random.nextInt(50) == 0) bytes(i) = 0 }
+ bytes
+ }
+
+ // Runs `compressRaw` into a buffer that is followed by sentinel bytes, and
returns whether the
+ // sentinel was overwritten, i.e. whether zstd wrote past the end of the
buffer.
+ private def writesPastEnd(
+ compressRaw: (BufferAllocator, ArrowBuf) => ArrowBuf,
+ data: Array[Byte]): Boolean = {
+ val guardSize = 64
+ val guardByte = 0x5A.toByte
+ var requestedSize = 0L
+ val guardedAllocator = new RootAllocator(Long.MaxValue) {
+ override def buffer(size: Long): ArrowBuf = {
+ requestedSize = size
+ val buf = super.buffer(size + guardSize)
+ buf.setBytes(size, Array.fill(guardSize)(guardByte))
+ buf
+ }
+ }
+ val inputAllocator =
+ ArrowUtils.rootAllocator.newChildAllocator(this.getClass.getSimpleName,
0, Long.MaxValue)
+ Utils.tryWithSafeFinally {
+ val input = inputAllocator.buffer(data.length)
+ val out = try {
+ input.setBytes(0, data)
+ input.writerIndex(data.length)
+ compressRaw(guardedAllocator, input)
+ } finally {
+ input.close()
+ }
+ try {
+ val guard = new Array[Byte](guardSize)
+ out.getBytes(requestedSize, guard)
+ !guard.forall(_ == guardByte)
+ } finally {
+ out.close()
+ }
+ } {
+
Utils.tryWithSafeFinally(guardedAllocator.close())(inputAllocator.close())
+ }
+ }
+
+ test("SPARK-60707: zstd compression does not write past the end of its
output buffer") {
+ val bounded = new RawBoundedZstdCompressionCodec
+ val stock = new RawStockZstdCompressionCodec
+ val random = new Random(0)
+ // Of these inputs, only those of roughly 280 to 1000 bytes make zstd fill
its whole
+ // destination; smaller and larger ones stop short of the end. The sizes
in that range are
+ // checked with many inputs each, and the other sizes only as smoke checks.
+ Seq(280, 500, 700).foreach { n =>
+ val inputs = Seq.fill(100)(noisyBytes(random, n))
+ inputs.foreach { data =>
+ assert(!writesPastEnd(bounded.compressRaw, data),
+ s"zstd wrote past the end of its output buffer (input size $n)")
+ }
+ // Control: these inputs 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.
+ assert(inputs.exists(writesPastEnd(stock.compressRaw, _)),
+ s"no input of size $n made Arrow's ZstdCompressionCodec write past its
output buffer")
+ }
+ Seq(1, 7, 100, 4096, 65536, 262144).foreach { n =>
+ assert(!writesPastEnd(bounded.compressRaw, noisyBytes(random, n)),
+ s"zstd wrote past the end of its output buffer (input size $n)")
+ }
+ }
+
+ 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)
+ val random = new Random(1)
+ Utils.tryWithSafeFinally {
+ // The write paths get the bounded codec, not Arrow's own.
+ assert(codec.isInstanceOf[BoundedZstdCompressionCodec])
+ Seq(1, 100, 4096, 100000).foreach { n =>
+ val compressible = Array.tabulate[Byte](n)(i => (i % 7).toByte)
+ val incompressible = random.nextBytes(n)
+ Seq(compressible, incompressible).foreach { expected =>
+ val input = allocator.buffer(n)
+ input.setBytes(0, expected)
+ input.writerIndex(n)
+ 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
Review Comment:
**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:
```scala
Seq(compressible -> (n >= 100), incompressible -> false).foreach {
case (expected, expectZstdFrame) =>
```
--
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]
