From 12aa61b79a21765734893f958625087841e8f508 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 28 Oct 2026 14:05:35 -0600 Subject: [PATCH 1/4] fix: give zstd the real destination capacity when compressing cached batches Arrow 18.3.0's ZstdCompressionCodec allocates compressBound(n) + 8 bytes, points zstd 8 bytes in past the length prefix, but still passes compressBound(n) + 8 as the capacity. zstd fills the destination up to the capacity it is given while entropy-coding a block that turns out not to compress, so the attempt can write up to 8 bytes past the buffer. The allocator's power-of-two rounding hides it except where the request lands within 8 bytes below a power of two, and without NMT the overrun is silent. Use a subclass whose doCompress passes compressBound(n), as the upstream fix (apache/arrow-java#1119, Arrow 20.0.0) does, and test the codec through an allocator that places a sentinel guard behind every buffer. Closes #7305 --- .../execution/arrow/CachedBatchIpc.scala | 9 +- .../arrow/CometZstdCompressionCodec.scala | 66 +++++ .../CometZstdCompressionCodecSuite.scala | 238 ++++++++++++++++++ 3 files changed, 309 insertions(+), 4 deletions(-) create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometZstdCompressionCodec.scala create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometZstdCompressionCodecSuite.scala 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..acdbd4d7ec1 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,10 @@ 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/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..15a70248498 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometZstdCompressionCodecSuite.scala @@ -0,0 +1,238 @@ +/* + * 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 + +import com.github.luben.zstd.Zstd + +/** + * 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. An overrun of the output buffer + * shows up as a changed sentinel instead of as silent corruption of whatever the allocator placed + * next, which is how it surfaces in production (#7305). The inputs are the ones that provoke it: + * bytes that look structured enough for zstd to try entropy coding but do not compress, in the + * sizes for which Arrow 18.3.0's codec asks for a capacity 8 bytes larger than the buffer. + */ +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"))) + + private val Prefix = CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH + + /** + * The input lengths whose compressed output buffer, as Arrow's codec sizes it, lands within 8 + * bytes below a power of two. Only there does the allocator's power-of-two rounding leave less + * slack than the 8 bytes by which Arrow overstates the capacity; everywhere else the rounding + * hides the overrun. 437 bytes, the size in #7305, is one of them: compressBound(437) + 8 is + * 509, allocated as 512. + */ + private val exposedLengths: Seq[Int] = (1 to 8192).filter { n => + val request = Prefix + Zstd.compressBound(n) + val rounded = java.lang.Long.highestOneBit(request - 1) << 1 + rounded - request < Prefix + } + + // 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) + input.setBytes(0, bytes) + input.writerIndex(bytes.length) + // 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) + exposedLengths.filter(_ < 1024).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 + exposedLengths.foreach { len => + (0 until 50).foreach { seed => + if (roundTrip(codec, root, guarded, protoLike(len, seed))) raw += 1 else packed += 1 + } + } + guarded.overruns shouldBe 0 + // The inputs are meant to defeat compression, which is the path that overruns. + raw should be > 0 + } + } + + test("Comet's zstd codec round-trips compressible and empty input") { + withAllocators { (root, guarded) => + val codec = new CometZstdCompressionCodec(3) + Seq(0, 1, 7, 437, 4096, 70000).foreach { len => + val storedRaw = roundTrip(codec, root, guarded, compressible(len, len)) + if (len > 16) storedRaw shouldBe false + } + guarded.overruns shouldBe 0 + } + } + + test("the guard notices a write past the buffer") { + withAllocators { (_, guarded) => + val buf = guarded.buffer(100) + // Comet runs with Arrow's bounds checks off, so a raw write is what an overrun looks like. + org.apache.arrow.memory.util.MemoryUtil.putByte(buf.memoryAddress() + 100, 1) + buf.close() + guarded.overruns shouldBe 1 + } + } +} + +/** + * 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) +} From c6418a7c15372bff9da95105787d153c4caefa02 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 28 Oct 2026 14:22:22 -0600 Subject: [PATCH 2/4] ci: run CometZstdCompressionCodecSuite in the PR build workflows --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + 2 files changed, 2 insertions(+) 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: | From 2c30c5762eafe5c81e0be853535bfa93cf5c3ed4 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 28 Oct 2026 14:41:44 -0600 Subject: [PATCH 3/4] address review: trim the zstd codec suite, run the cache suite through Comet's codec --- .../execution/arrow/CachedBatchIpc.scala | 3 +- .../comet/exec/CometInMemoryCacheSuite.scala | 85 +++++++++---------- .../CometZstdCompressionCodecSuite.scala | 64 ++++---------- 3 files changed, 60 insertions(+), 92 deletions(-) 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 acdbd4d7ec1..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 @@ -72,7 +72,8 @@ private[comet] object CachedBatchIpc { case "none" => NoCompressionCodec.INSTANCE // 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. + // 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. 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 index 15a70248498..e4f0e07afd0 100644 --- 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 @@ -32,17 +32,14 @@ import org.apache.arrow.memory.{AllocationListener, AllocationReservation, Arrow import org.apache.arrow.memory.rounding.RoundingPolicy import org.apache.arrow.vector.compression.CompressionUtil -import com.github.luben.zstd.Zstd - /** * 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. An overrun of the output buffer - * shows up as a changed sentinel instead of as silent corruption of whatever the allocator placed - * next, which is how it surfaces in production (#7305). The inputs are the ones that provoke it: - * bytes that look structured enough for zstd to try entropy coding but do not compress, in the - * sizes for which Arrow 18.3.0's codec asks for a capacity 8 bytes larger than the buffer. + * 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 { @@ -50,20 +47,9 @@ class CometZstdCompressionCodecSuite extends AnyFunSuite with Matchers { // target/tmp. Nothing else creates that directory in a fresh checkout. Files.createDirectories(Paths.get(System.getProperty("java.io.tmpdir"))) - private val Prefix = CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH - - /** - * The input lengths whose compressed output buffer, as Arrow's codec sizes it, lands within 8 - * bytes below a power of two. Only there does the allocator's power-of-two rounding leave less - * slack than the 8 bytes by which Arrow overstates the capacity; everywhere else the rounding - * hides the overrun. 437 bytes, the size in #7305, is one of them: compressBound(437) + 8 is - * 509, allocated as 512. - */ - private val exposedLengths: Seq[Int] = (1 to 8192).filter { n => - val request = Prefix + Zstd.compressBound(n) - val rounded = java.lang.Long.highestOneBit(request - 1) << 1 - rounded - request < Prefix - } + // 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 @@ -120,7 +106,7 @@ class CometZstdCompressionCodecSuite extends AnyFunSuite with Matchers { // go together. withAllocators { (root, guarded) => val codec = new ZstdCompressionCodec(1) - exposedLengths.filter(_ < 1024).foreach { len => + Lengths.foreach { len => (0 until 50).foreach(seed => roundTrip(codec, root, guarded, protoLike(len, seed))) } guarded.overruns should be > 0 @@ -132,35 +118,17 @@ class CometZstdCompressionCodecSuite extends AnyFunSuite with Matchers { val codec = new CometZstdCompressionCodec(1) var raw = 0 var packed = 0 - exposedLengths.foreach { len => - (0 until 50).foreach { seed => - if (roundTrip(codec, root, guarded, protoLike(len, seed))) raw += 1 else packed += 1 - } + 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 - // The inputs are meant to defeat compression, which is the path that overruns. + // Both outcomes ran under the guard: the stored block the overrun comes from, and a block + // that really compressed. raw should be > 0 - } - } - - test("Comet's zstd codec round-trips compressible and empty input") { - withAllocators { (root, guarded) => - val codec = new CometZstdCompressionCodec(3) - Seq(0, 1, 7, 437, 4096, 70000).foreach { len => - val storedRaw = roundTrip(codec, root, guarded, compressible(len, len)) - if (len > 16) storedRaw shouldBe false - } - guarded.overruns shouldBe 0 - } - } - - test("the guard notices a write past the buffer") { - withAllocators { (_, guarded) => - val buf = guarded.buffer(100) - // Comet runs with Arrow's bounds checks off, so a raw write is what an overrun looks like. - org.apache.arrow.memory.util.MemoryUtil.putByte(buf.memoryAddress() + 100, 1) - buf.close() - guarded.overruns shouldBe 1 + packed should be > 0 } } } From 03ce3b9f00140f6cbeccbc41725c80fd2e1679c3 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 28 Oct 2026 16:06:08 -0600 Subject: [PATCH 4/4] test: widen Int lengths explicitly in the zstd codec suite --- .../execution/arrow/CometZstdCompressionCodecSuite.scala | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 index e4f0e07afd0..87463d1d70b 100644 --- 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 @@ -72,9 +72,9 @@ class CometZstdCompressionCodecSuite extends AnyFunSuite with Matchers { allocator: BufferAllocator, guarded: GuardedAllocator, bytes: Array[Byte]): Boolean = { - val input = allocator.buffer(bytes.length) + val input = allocator.buffer(bytes.length.toLong) input.setBytes(0, bytes) - input.writerIndex(bytes.length) + 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