andygrove commented on code in PR #59922:
URL: https://github.com/apache/spark/pull/59922#discussion_r4428793690
##########
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:
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).
##########
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:
Added in 3ab13939845: 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.
##########
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:
Reworded in the description as suggested.
##########
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:
Thanks for checking. I'll keep `internalError`; the error name stays in the
message, so a native allocation failure can still be diagnosed.
##########
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:
Done in 3ab13939845, using your pairing.
--
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]
