dongjoon-hyun commented on code in PR #59922:
URL: https://github.com/apache/spark/pull/59922#discussion_r4416019839
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtilsSuite.scala:
##########
@@ -85,4 +93,75 @@ class ArrowCompressionUtilsSuite extends SparkFunSuite {
}
assert(e.getMessage.contains("Unsupported Arrow compression codec:
snappy"))
}
+
+ test("SPARK-60707: zstd compression does not write past the end of its
output buffer") {
+ val guardSize = 64
+ val guardByte = 0x5A.toByte
+ // Pads each output buffer with sentinel bytes so that an overrun is
detectable.
+ 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)
+ val codec = new RawBoundedZstdCompressionCodec
+ 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 =>
+ val input = inputAllocator.buffer(n)
+ val out = try {
+ input.setBytes(0, random.nextBytes(n))
+ input.writerIndex(n)
+ codec.compressRaw(guardedAllocator, input)
+ } finally {
+ input.close()
+ }
+ try {
+ val guard = new Array[Byte](guardSize)
+ out.getBytes(requestedSize, guard)
+ assert(guard.forall(_ == guardByte),
+ s"zstd wrote past the end of a $requestedSize byte output buffer
(input size $n)")
+ } finally {
+ out.close()
+ }
+ }
+ } finally {
+ guardedAllocator.close()
+ inputAllocator.close()
+ }
+ }
+
+ 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)
+ try {
+ Seq(1, 100, 4096, 100000).foreach { n =>
+ val compressible = Array.tabulate[Byte](n)(i => (i % 7).toByte)
+ val incompressible = random.nextBytes(n)
Review Comment:
**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?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtilsSuite.scala:
##########
@@ -85,4 +93,75 @@ class ArrowCompressionUtilsSuite extends SparkFunSuite {
}
assert(e.getMessage.contains("Unsupported Arrow compression codec:
snappy"))
}
+
+ test("SPARK-60707: zstd compression does not write past the end of its
output buffer") {
+ val guardSize = 64
+ val guardByte = 0x5A.toByte
+ // Pads each output buffer with sentinel bytes so that an overrun is
detectable.
+ 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)
+ val codec = new RawBoundedZstdCompressionCodec
+ 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 =>
Review Comment:
**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.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala:
##########
@@ -267,18 +266,7 @@ private object ArrowCachedBatchSerializer {
def createCompressionCodec(
codecName: String,
compressionLevel: Int): CompressionCodec = {
- codecName.toLowerCase match {
- case "none" => NoCompressionCodec.INSTANCE
- // The codec instance must be constructed directly so that
compressionLevel is honored:
- // CompressionCodec.Factory.createCodec(codecType) ignores the level and
builds a codec at
- // the default level. The level only matters on the write side; the read
side looks up the
- // codec by the type recorded in the IPC message.
- case "zstd" => new ZstdCompressionCodec(compressionLevel)
- case "lz4" => new Lz4CompressionCodec()
- case other =>
- throw SparkException.internalError(
- s"Unsupported Arrow compression codec: $other. Supported values:
none, zstd, lz4")
- }
+ ArrowCompressionUtils.createCompressionCodec(codecName.toLowerCase,
compressionLevel)
Review Comment:
**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.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtilsSuite.scala:
##########
@@ -85,4 +93,75 @@ class ArrowCompressionUtilsSuite extends SparkFunSuite {
}
assert(e.getMessage.contains("Unsupported Arrow compression codec:
snappy"))
}
+
+ test("SPARK-60707: zstd compression does not write past the end of its
output buffer") {
+ val guardSize = 64
+ val guardByte = 0x5A.toByte
+ // Pads each output buffer with sentinel bytes so that an overrun is
detectable.
+ 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)
+ val codec = new RawBoundedZstdCompressionCodec
+ 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 =>
+ val input = inputAllocator.buffer(n)
+ val out = try {
+ input.setBytes(0, random.nextBytes(n))
+ input.writerIndex(n)
+ codec.compressRaw(guardedAllocator, input)
+ } finally {
+ input.close()
+ }
+ try {
+ val guard = new Array[Byte](guardSize)
+ out.getBytes(requestedSize, guard)
+ assert(guard.forall(_ == guardByte),
+ s"zstd wrote past the end of a $requestedSize byte output buffer
(input size $n)")
+ } finally {
+ out.close()
+ }
+ }
+ } finally {
+ guardedAllocator.close()
+ inputAllocator.close()
+ }
+ }
+
+ 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)
+ try {
+ 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 decompressed = codec.decompress(allocator,
codec.compress(allocator, input))
+ try {
+ val actual = new Array[Byte](n)
+ decompressed.getBytes(0, actual)
+ assert(actual.sameElements(expected))
+ } finally {
+ decompressed.close()
+ }
+ }
+ }
+ } finally {
+ allocator.close()
Review Comment:
**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.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtils.scala:
##########
@@ -40,11 +42,41 @@ private[sql] object ArrowCompressionUtils {
zstdCompressionLevel: Int): CompressionCodec = {
codecName match {
case "none" => NoCompressionCodec.INSTANCE
- case "zstd" => new ZstdCompressionCodec(zstdCompressionLevel)
+ case "zstd" => new BoundedZstdCompressionCodec(zstdCompressionLevel)
Review Comment:
**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](https://github.com/apache/spark/pull/59922#discussion_r4414638639)
would catch that in tests, but the scaladoc would still suggest that only the
level matters. Could it add a sentence such as:
```scala
* For zstd it must also be a [[BoundedZstdCompressionCodec]], which
neither factory overload
* returns, not even `createCodec(codecType, compressionLevel)`, which
does keep the level.
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtils.scala:
##########
@@ -40,11 +42,41 @@ private[sql] object ArrowCompressionUtils {
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 18.3 and
Review Comment:
**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.
--
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]
