fix: give zstd the real destination capacity when compressing cached batches

.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: |

.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: |

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

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
+  }
+}

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 {