peter-toth commented on code in PR #59922:
URL: https://github.com/apache/spark/pull/59922#discussion_r4414638639
##########
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:
**Finding 1.** The guard test above builds `RawBoundedZstdCompressionCodec`
directly, and this round-trip passes with either codec. I switched
`createCompressionCodec` back to `new
ZstdCompressionCodec(zstdCompressionLevel)`, and `ArrowCompressionUtilsSuite`
and `ArrowCachedBatchSerializerSuite` still passed (83 tests). With this check,
the round-trip test fails on that change and passes here.
`ArrowCachedBatchSerializer` now goes through the same factory, so it covers
that path too:
```suggestion
val codec = ArrowCompressionUtils.createCompressionCodec("zstd", 3)
// The write paths get the bounded codec, not Arrow's own.
assert(codec.isInstanceOf[BoundedZstdCompressionCodec])
```
--
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]
