dongjoon-hyun commented on code in PR #59922:
URL: https://github.com/apache/spark/pull/59922#discussion_r4418702347
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtilsSuite.scala:
##########
@@ -85,4 +101,110 @@ class ArrowCompressionUtilsSuite extends SparkFunSuite {
}
assert(e.getMessage.contains("Unsupported Arrow compression codec:
snappy"))
}
+
+ // 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)
+ try {
+ 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()
+ }
+ } finally {
+ 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)
+ Seq(1, 7, 100, 280, 1000, 4096, 65536, 262144).foreach { n =>
Review Comment:
**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?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtilsSuite.scala:
##########
@@ -85,4 +101,110 @@ class ArrowCompressionUtilsSuite extends SparkFunSuite {
}
assert(e.getMessage.contains("Unsupported Arrow compression codec:
snappy"))
}
+
+ // 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)
+ try {
+ 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()
+ }
+ } finally {
+ 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)
+ Seq(1, 7, 100, 280, 1000, 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)")
+ }
+ // At 500 bytes a large share of these inputs make Arrow's own codec
overrun the buffer.
+ val inputs = Seq.fill(100)(noisyBytes(random, 500))
+ inputs.foreach { data =>
+ assert(!writesPastEnd(bounded.compressRaw, data),
+ "zstd wrote past the end of its output buffer (input size 500)")
+ }
+ // 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.
Review Comment:
**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:
```scala
// 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.
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala:
##########
@@ -1005,7 +985,7 @@ private class InternalRowToArrowCachedBatchIterator(
compressionCodecName: String,
compressionLevel: Int) extends Iterator[ArrowCachedBatch] {
- private val compressionCodec =
ArrowCachedBatchSerializer.createCompressionCodec(
+ private val compressionCodec = ArrowCompressionUtils.createCompressionCodec(
Review Comment:
**3. [Low] The PR description still describes the first commit**
Since 23dc6e39c9f, 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.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtilsSuite.scala:
##########
@@ -85,4 +101,110 @@ class ArrowCompressionUtilsSuite extends SparkFunSuite {
}
assert(e.getMessage.contains("Unsupported Arrow compression codec:
snappy"))
}
+
+ // 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)
+ try {
+ 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()
+ }
+ } finally {
+ guardedAllocator.close()
+ inputAllocator.close()
+ }
Review Comment:
**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:
```scala
Utils.tryWithSafeFinally {
...
} {
Utils.tryWithSafeFinally(guardedAllocator.close())(inputAllocator.close())
}
```
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtilsSuite.scala:
##########
@@ -85,4 +101,110 @@ class ArrowCompressionUtilsSuite extends SparkFunSuite {
}
assert(e.getMessage.contains("Unsupported Arrow compression codec:
snappy"))
}
+
+ // 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)
+ try {
+ 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()
+ }
+ } finally {
+ 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)
+ Seq(1, 7, 100, 280, 1000, 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)")
+ }
+ // At 500 bytes a large share of these inputs make Arrow's own codec
overrun the buffer.
+ val inputs = Seq.fill(100)(noisyBytes(random, 500))
+ inputs.foreach { data =>
+ assert(!writesPastEnd(bounded.compressRaw, data),
+ "zstd wrote past the end of its output buffer (input size 500)")
+ }
+ // 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.
+ assert(inputs.exists(writesPastEnd(stock.compressRaw, _)),
+ "none of the inputs made Arrow's ZstdCompressionCodec write past its
output buffer")
+ }
+
+ 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
+ assert(compressed.getLong(0) === (if (expectZstdFrame) n.toLong else
-1L))
Review Comment:
**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`):
```scala
val prefix = compressed.nioBuffer(0,
8).order(ByteOrder.LITTLE_ENDIAN).getLong(0)
assert(prefix === (if (expectZstdFrame) n.toLong else -1L))
```
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtilsSuite.scala:
##########
@@ -28,9 +32,21 @@ import org.apache.arrow.vector.ipc.message.MessageSerializer
import org.apache.spark.{SparkException, SparkFunSuite}
import org.apache.spark.sql.types.{StringType, StructField, StructType}
import org.apache.spark.sql.util.ArrowUtils
+import org.apache.spark.util.Utils
class ArrowCompressionUtilsSuite extends SparkFunSuite {
+ // 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)
+ }
Review Comment:
**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.
##########
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 new RuntimeException("Error compressing: " +
Zstd.getErrorName(bytesWritten))
Review Comment:
**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.
```suggestion
throw SparkException.internalError("Error compressing: " +
Zstd.getErrorName(bytesWritten))
```
--
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]
