andygrove commented on code in PR #59922:
URL: https://github.com/apache/spark/pull/59922#discussion_r4418467997
##########
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:
Done in 23dc6e39c9f: the round-trip test uses `Utils.tryWithSafeFinally {
... } { allocator.close() }` so the first error is kept.
##########
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:
Done in 23dc6e39c9f: the scaladoc now says "arrow-java up to 19.0". I
updated the PR description to match.
--
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]
