andygrove commented on code in PR #59922:
URL: https://github.com/apache/spark/pull/59922#discussion_r4418466456
##########
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)
Review Comment:
Done in 23dc6e39c9f: the round-trip test now asserts
`codec.isInstanceOf[BoundedZstdCompressionCodec]`. Thanks!
##########
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:
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 23dc6e39c9f.
##########
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:
Done in 23dc6e39c9f: 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.
##########
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:
Done in 23dc6e39c9f: added the sentence about `BoundedZstdCompressionCodec`
to the scaladoc, including that `createCodec(codecType, compressionLevel)`
keeps the level but still returns the stock codec.
##########
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:
Agreed, and I checked that `SQLConf` lowercases and validates the value.
Done in 23dc6e39c9f: removed
`ArrowCachedBatchSerializer.createCompressionCodec`, both callers use
`ArrowCompressionUtils.createCompressionCodec` directly, and the duplicate
level test is gone along with the now-unused imports.
--
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]
