andygrove commented on code in PR #59922:
URL: https://github.com/apache/spark/pull/59922#discussion_r4423969670
##########
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:
Thanks for the table. Done in 1464f9cfcb2: the main loop now uses 280, 500
and 700 bytes with 100 inputs each, and the comment says only roughly 280 to
1000 bytes can reach the end of the buffer. The other sizes (1, 7, 100, 4096,
65536, 262144) stay as smoke checks and the comment says so.
##########
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:
Done in 1464f9cfcb2: the control is now per size, and its comment names both
causes (Arrow upgrade: remove the class; zstd-jni upgrade: pick inputs that
still overrun), using your wording.
##########
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:
Updated the description: the removed wrapper and its duplicate test, the new
guard test and controls, the round-trip assertions, and the local run (82 tests
across both suites).
##########
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:
Done in 1464f9cfcb2, using the nested `Utils.tryWithSafeFinally` from your
suggestion.
##########
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:
Done in 1464f9cfcb2: the prefix is read through `nioBuffer(0,
8).order(ByteOrder.LITTLE_ENDIAN)`.
--
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]
