andygrove commented on code in PR #59922:
URL: https://github.com/apache/spark/pull/59922#discussion_r4423971141
##########
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 new RuntimeException("Error compressing: " +
Zstd.getErrorName(bytesWritten))
Review Comment:
Done in 1464f9cfcb2: now `SparkException.internalError`, which fits the rest
of the file. With `compressBound` as the capacity zstd cannot run out of room,
so an error here is unexpected.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtilsSuite.scala:
##########
@@ -28,9 +32,21 @@ import org.apache.arrow.vector.ipc.message.MessageSerializer
import org.apache.spark.{SparkException, SparkFunSuite}
import org.apache.spark.sql.types.{StringType, StructField, StructType}
import org.apache.spark.sql.util.ArrowUtils
+import org.apache.spark.util.Utils
class ArrowCompressionUtilsSuite extends SparkFunSuite {
+ // 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)
+ }
Review Comment:
Done in 1464f9cfcb2: both `Raw*` classes now sit with `noisyBytes` and
`writesPastEnd` above the guard test.
--
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]
