diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 52e86f59ffd..a35444eb25b 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -633,6 +633,7 @@ jobs: org.apache.spark.sql.comet.execution.arrow.CachedBatchRowIteratorSuite org.apache.spark.sql.comet.execution.arrow.CometArrowWriterSuite org.apache.spark.sql.comet.execution.arrow.CometStringWriterSuite + org.apache.spark.sql.comet.execution.arrow.CometZstdCompressionCodecSuite org.apache.spark.sql.CometSparkInternalFunctionsSuite - name: "expressions" value: | diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index b1d89ec1240..9547ece6e4d 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -299,6 +299,7 @@ jobs: org.apache.spark.sql.comet.execution.arrow.CachedBatchRowIteratorSuite org.apache.spark.sql.comet.execution.arrow.CometArrowWriterSuite org.apache.spark.sql.comet.execution.arrow.CometStringWriterSuite + org.apache.spark.sql.comet.execution.arrow.CometZstdCompressionCodecSuite org.apache.spark.sql.CometSparkInternalFunctionsSuite - name: "expressions" value: | diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala index 81f8e56918a..80bdd4109e0 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala @@ -26,7 +26,7 @@ import scala.collection.mutable import scala.jdk.CollectionConverters._ import scala.util.control.NonFatal -import org.apache.arrow.compression.{CommonsCompressionFactory, ZstdCompressionCodec} +import org.apache.arrow.compression.CommonsCompressionFactory import org.apache.arrow.flatbuf.{RecordBatch => FlatBufRecordBatch} import org.apache.arrow.memory.{ArrowBuf, BufferAllocator} import org.apache.arrow.vector.{FieldVector, TypeLayout, ValueVector, VectorLoader, VectorSchemaRoot, VectorUnloader} @@ -70,9 +70,11 @@ private[comet] object CachedBatchIpc { */ def compressionCodec(codecName: String, zstdLevel: Int): CompressionCodec = codecName match { case "none" => NoCompressionCodec.INSTANCE - // Constructed directly rather than through CompressionCodec.Factory, which ignores the level - // and always builds a codec at zstd's default. - case "zstd" => new ZstdCompressionCodec(zstdLevel) + // Comet's subclass of Arrow's codec, which corrects the capacity handed to zstd; see + // CometZstdCompressionCodec. Constructed directly rather than through + // CompressionCodec.Factory, which ignores the level and always builds a codec at zstd's + // default. + case "zstd" => new CometZstdCompressionCodec(zstdLevel) // Arrow's LZ4_FRAME format, written with lz4-java rather than Arrow's own pure-Java codec; // see Lz4FrameCompressionCodec. case "lz4" => Lz4FrameCompressionCodec diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometZstdCompressionCodec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometZstdCompressionCodec.scala new file mode 100644 index 00000000000..b6f0c556c5e --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometZstdCompressionCodec.scala @@ -0,0 +1,66 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.comet.execution.arrow + +import org.apache.arrow.compression.ZstdCompressionCodec +import org.apache.arrow.memory.{ArrowBuf, BufferAllocator} +import org.apache.arrow.vector.compression.CompressionUtil +import org.apache.spark.SparkException + +import com.github.luben.zstd.Zstd + +/** + * Arrow's zstd codec, with the destination capacity it hands zstd corrected. + * + * Arrow 18.3.0's `ZstdCompressionCodec.doCompress` allocates `compressBound(n) + 8` bytes and + * points zstd 8 bytes in, past the uncompressed-length prefix, but still tells it the capacity is + * `compressBound(n) + 8`. What zstd guarantees is that it never writes past the capacity it is + * given, not that it never writes past `compressBound`: while it entropy-codes the literals of a + * block that turns out not to compress, it fills the destination up to that capacity before it + * falls back to a stored block, so the attempt can run up to 8 bytes past the buffer. The + * allocator's power-of-two rounding absorbs that unless the request lands within 8 bytes below a + * power of two, and when it does, the overrun lands in whatever follows the buffer without any + * error. With Native Memory Tracking on, the JVM aborts when the buffer is freed (#7305). + * + * Fixed upstream by apache/arrow-java#1119 for Arrow 20.0.0; this is that fix. Only `doCompress` + * differs, so decompression stays Arrow's. + */ +private[comet] class CometZstdCompressionCodec(level: Int) extends ZstdCompressionCodec(level) { + + override protected def doCompress( + allocator: BufferAllocator, + uncompressed: ArrowBuf): ArrowBuf = { + val capacity = Zstd.compressBound(uncompressed.writerIndex()) + val compressed = allocator.buffer(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH + capacity) + val written = Zstd.compressUnsafe( + compressed.memoryAddress() + CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH, + capacity, + uncompressed.memoryAddress(), + uncompressed.writerIndex(), + level) + if (Zstd.isError(written)) { + compressed.close() + throw new SparkException( + s"zstd failed to compress a cached batch: ${Zstd.getErrorName(written)}") + } + compressed.writerIndex(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH + written) + compressed + } +} diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 2e1f4d56d33..2f55cd93ee5 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -28,7 +28,6 @@ import scala.concurrent.ExecutionContext.Implicits.global import scala.concurrent.duration.DurationInt import scala.jdk.CollectionConverters._ -import org.apache.arrow.compression.ZstdCompressionCodec import org.apache.arrow.memory.{ArrowBuf, BufferAllocator} import org.apache.arrow.vector.{BitVector, FixedSizeBinaryVector, IntVector, LargeVarCharVector, VarBinaryVector, VarCharVector} import org.apache.arrow.vector.compression.{CompressionCodec, CompressionUtil, NoCompressionCodec} @@ -2616,49 +2615,49 @@ class CometInMemoryCacheSuite extends CometTestBase { // uncompressed-length prefix of zero, and without a codec each is a zero-length window, so // this is the batch where every window the read path copies out is empty or nearly so. val cacheSchema = StructType(Seq(StructField("i", IntegerType), StructField("s", StringType))) - Seq[CompressionCodec](NoCompressionCodec.INSTANCE, new ZstdCompressionCodec(1)).foreach { - codec => - val ints = new IntVector("i", CometArrowAllocator) - val strings = new VarCharVector("s", CometArrowAllocator) - val cached = + val zstd = CometCachedBatchHelper.compressionCodec("zstd") + Seq[CompressionCodec](NoCompressionCodec.INSTANCE, zstd).foreach { codec => + val ints = new IntVector("i", CometArrowAllocator) + val strings = new VarCharVector("s", CometArrowAllocator) + val cached = + try { + ints.allocateNew() + ints.setValueCount(0) + strings.allocateNew() + strings.setValueCount(0) + val batch = new ColumnarBatch( + Array[ColumnVector](new CometPlainVector(ints), new CometPlainVector(strings)), + 0) + CometCachedBatchHelper.cachedBatch( + CometCachedBatchHelper.serialize(batch, codec, CometArrowAllocator), + 0) + } finally { + ints.close() + strings.close() + } + + val allocator = + CometArrowAllocator.newChildAllocator( + s"zero-rows-${codec.getCodecType}", + 0, + Long.MaxValue) + try { + Seq(Array(0, 1), Array(1), Array(0)).foreach { selected => + val root = CometCachedBatchHelper.load(cached, cacheSchema, selected, allocator) try { - ints.allocateNew() - ints.setValueCount(0) - strings.allocateNew() - strings.setValueCount(0) - val batch = new ColumnarBatch( - Array[ColumnVector](new CometPlainVector(ints), new CometPlainVector(strings)), - 0) - CometCachedBatchHelper.cachedBatch( - CometCachedBatchHelper.serialize(batch, codec, CometArrowAllocator), - 0) + assert(root.getRowCount == 0) + assert( + root.getSchema.getFields.asScala.map(_.getName) == + selected.map(cacheSchema(_).name).toSeq, + s"${codec.getCodecType} read the wrong columns for ${selected.mkString(",")}") } finally { - ints.close() - strings.close() + root.close() } - - val allocator = - CometArrowAllocator.newChildAllocator( - s"zero-rows-${codec.getCodecType}", - 0, - Long.MaxValue) - try { - Seq(Array(0, 1), Array(1), Array(0)).foreach { selected => - val root = CometCachedBatchHelper.load(cached, cacheSchema, selected, allocator) - try { - assert(root.getRowCount == 0) - assert( - root.getSchema.getFields.asScala.map(_.getName) == - selected.map(cacheSchema(_).name).toSeq, - s"${codec.getCodecType} read the wrong columns for ${selected.mkString(",")}") - } finally { - root.close() - } - } - assert(allocator.getAllocatedMemory == 0) - } finally { - allocator.close() } + assert(allocator.getAllocatedMemory == 0) + } finally { + allocator.close() + } } } @@ -3207,7 +3206,7 @@ class CometInMemoryCacheSuite extends CometTestBase { * failure -- zstd unable to allocate its workspace part way through a batch -- leaves behind. */ private class FailAfterCompressionCodec(succeedFor: Int) extends CompressionCodec { - private val delegate = new ZstdCompressionCodec(1) + private val delegate = CometCachedBatchHelper.compressionCodec("zstd") var compressed: Int = 0 override def compress(allocator: BufferAllocator, buffer: ArrowBuf): ArrowBuf = { @@ -3315,7 +3314,7 @@ class CometInMemoryCacheSuite extends CometTestBase { failure } - assert(write(new ZstdCompressionCodec(1)).isEmpty) + assert(write(CometCachedBatchHelper.compressionCodec("zstd")).isEmpty) // The int column's two buffers compress, then the copy's first one fails. val codec = new FailAfterCompressionCodec(succeedFor = 2) @@ -3357,7 +3356,7 @@ class CometInMemoryCacheSuite extends CometTestBase { CometCachedBatchHelper.cachedBatch( CometCachedBatchHelper.serialize( batch, - new ZstdCompressionCodec(1), + CometCachedBatchHelper.compressionCodec("zstd"), CometArrowAllocator), rows) } finally { diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometZstdCompressionCodecSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometZstdCompressionCodecSuite.scala new file mode 100644 index 00000000000..87463d1d70b --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometZstdCompressionCodecSuite.scala @@ -0,0 +1,206 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.comet.execution.arrow + +import java.nio.file.{Files, Paths} +import java.util.Random + +import scala.collection.mutable + +import org.scalatest.funsuite.AnyFunSuite +import org.scalatest.matchers.should.Matchers + +import org.apache.arrow.compression.ZstdCompressionCodec +import org.apache.arrow.memory.{AllocationListener, AllocationReservation, ArrowBuf, BufferAllocator, BufferManager, ForeignAllocation, RootAllocator} +import org.apache.arrow.memory.rounding.RoundingPolicy +import org.apache.arrow.vector.compression.CompressionUtil + +/** + * Checks that the zstd codec never writes past the buffer it allocates for its output. + * + * The buffers come from [[GuardedAllocator]], which hands the codec exactly the capacity it asked + * for and keeps a sentinel-filled guard region right behind it, so an overrun shows up as a + * changed sentinel instead of as silent corruption of whatever the allocator placed next (#7305). + * The defect is explained on [[CometZstdCompressionCodec]]. The inputs are the ones that provoke + * it: bytes that look structured enough for zstd to try entropy coding but do not compress. + */ +class CometZstdCompressionCodecSuite extends AnyFunSuite with Matchers { + + // zstd-jni unpacks its native library into java.io.tmpdir, which the test JVM points at + // target/tmp. Nothing else creates that directory in a fresh checkout. + Files.createDirectories(Paths.get(System.getProperty("java.io.tmpdir"))) + + // The guard sees the overrun at any length, since the allocator cuts each buffer to the exact + // size requested. This range just gives it many sizes to look at. + private val Lengths = 400 to 600 + + // Protobuf-like bytes: random payload with a scatter of repeated field tags. The tags make the + // block pass zstd's "worth trying" heuristic, the payload makes the attempt not pay off, so + // zstd runs its entropy coder to the end of the capacity it was given before storing the block. + private def protoLike(len: Int, seed: Int): Array[Byte] = { + val rnd = new Random(seed * 7919L + len) + val out = new Array[Byte](len) + rnd.nextBytes(out) + (0 until 8 + rnd.nextInt(24)).foreach { _ => + out(rnd.nextInt(len)) = (0x08 + 8 * rnd.nextInt(4)).toByte + } + out + } + + private def compressible(len: Int, seed: Int): Array[Byte] = + Array.tabulate(len)(i => ((i / 7 + seed) % 5).toByte) + + private def roundTrip( + codec: ZstdCompressionCodec, + allocator: BufferAllocator, + guarded: GuardedAllocator, + bytes: Array[Byte]): Boolean = { + val input = allocator.buffer(bytes.length.toLong) + input.setBytes(0, bytes) + input.writerIndex(bytes.length.toLong) + // compress consumes the input reference and allocates its output from the guarded allocator. + val compressed = codec.compress(guarded, input) + val storedRaw = compressed.getLong(0) == CompressionUtil.NO_COMPRESSION_LENGTH + val plain = codec.decompress(allocator, compressed) + try { + plain.writerIndex() shouldBe bytes.length + val actual = new Array[Byte](bytes.length) + plain.getBytes(0, actual) + actual shouldBe bytes + } finally { + plain.close() + } + storedRaw + } + + private def withAllocators(f: (BufferAllocator, GuardedAllocator) => Unit): Unit = { + val root = new RootAllocator() + val guarded = new GuardedAllocator(root) + try f(root, guarded) + finally { + guarded.close() + root.close() + } + } + + test("Arrow's zstd codec writes past the output buffer it allocates") { + // Documents the defect this suite guards against (apache/arrow-java#1119, fixed in Arrow + // 20.0.0). Once Comet moves to an Arrow with the fix, CometZstdCompressionCodec and this test + // go together. + withAllocators { (root, guarded) => + val codec = new ZstdCompressionCodec(1) + Lengths.foreach { len => + (0 until 50).foreach(seed => roundTrip(codec, root, guarded, protoLike(len, seed))) + } + guarded.overruns should be > 0 + } + } + + test("Comet's zstd codec stays within the output buffer it allocates") { + withAllocators { (root, guarded) => + val codec = new CometZstdCompressionCodec(1) + var raw = 0 + var packed = 0 + def count(storedRaw: Boolean): Unit = if (storedRaw) raw += 1 else packed += 1 + Lengths.foreach { len => + (0 until 50).foreach(seed => count(roundTrip(codec, root, guarded, protoLike(len, seed)))) + (0 until 5).foreach(seed => + count(roundTrip(codec, root, guarded, compressible(len, seed)))) + } + guarded.overruns shouldBe 0 + // Both outcomes ran under the guard: the stored block the overrun comes from, and a block + // that really compressed. + raw should be > 0 + packed should be > 0 + } + } +} + +/** + * Hands out buffers of exactly the requested capacity, each followed by a sentinel-filled guard + * region, and counts the buffers whose guard was written to by the time they were released. + * + * Each buffer is a slice at the start of a larger allocation from `root`; the allocator keeps a + * reference to the allocation so the guard can still be read once the codec has closed its slice. + * Everything other than allocation delegates to `root`. + */ +private final class GuardedAllocator(root: BufferAllocator) extends BufferAllocator { + private val Guard = 64 + private val Sentinel: Byte = 0x5a + // Each allocation with the size that was requested from it, where its guard begins. + private val parents = mutable.ArrayBuffer.empty[(ArrowBuf, Long)] + + /** Buffers whose guard region no longer holds the sentinel. */ + def overruns: Int = parents.count { case (parent, size) => + (0 until Guard).exists(i => parent.getByte(size + i) != Sentinel) + } + + override def buffer(size: Long): ArrowBuf = { + val parent = root.buffer(size + Guard) + // The slice is cut to exactly `size`, whatever the allocator rounded the request to, and the + // guard starts where the slice ends. + parent.setBytes(size, Array.fill[Byte](Guard)(Sentinel)) + // The slice shares the parent's reference manager, so retain once for the reference this + // allocator keeps; the codec's close releases the other. + parent.getReferenceManager.retain() + parents += ((parent, size)) + parent.slice(0, size).writerIndex(0) + } + + override def buffer(size: Long, manager: BufferManager): ArrowBuf = buffer(size) + + override def close(): Unit = { + parents.foreach(_._1.getReferenceManager.release()) + parents.clear() + } + + override def getRoot: BufferAllocator = root.getRoot + override def newChildAllocator( + name: String, + initReservation: Long, + maxAllocation: Long): BufferAllocator = + root.newChildAllocator(name, initReservation, maxAllocation) + override def newChildAllocator( + name: String, + listener: AllocationListener, + initReservation: Long, + maxAllocation: Long): BufferAllocator = + root.newChildAllocator(name, listener, initReservation, maxAllocation) + override def getAllocatedMemory: Long = root.getAllocatedMemory + override def getLimit: Long = root.getLimit + override def getInitReservation: Long = root.getInitReservation + override def setLimit(newLimit: Long): Unit = root.setLimit(newLimit) + override def getPeakMemoryAllocation: Long = root.getPeakMemoryAllocation + override def getHeadroom: Long = root.getHeadroom + override def forceAllocate(size: Long): Boolean = root.forceAllocate(size) + override def releaseBytes(size: Long): Unit = root.releaseBytes(size) + override def getListener: AllocationListener = root.getListener + override def getParentAllocator: BufferAllocator = root.getParentAllocator + override def getChildAllocators: java.util.Collection[BufferAllocator] = root.getChildAllocators + override def newReservation(): AllocationReservation = root.newReservation() + override def getEmpty: ArrowBuf = root.getEmpty + override def getName: String = "guarded" + override def isOverLimit: Boolean = root.isOverLimit + override def toVerboseString: String = root.toVerboseString + override def assertOpen(): Unit = root.assertOpen() + override def getRoundingPolicy: RoundingPolicy = root.getRoundingPolicy + override def wrapForeignAllocation(allocation: ForeignAllocation): ArrowBuf = + root.wrapForeignAllocation(allocation) +}