[SPARK-60776][SQL] Fix Arrow memory leak when compressing a record batch fails

sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtils.scala

@@ -17,10 +17,16 @@
 
 package org.apache.spark.sql.execution.arrow
 
+import scala.collection.mutable.ArrayBuffer
+import scala.jdk.CollectionConverters._
+
 import com.github.luben.zstd.Zstd
 import org.apache.arrow.compression.{Lz4CompressionCodec, ZstdCompressionCodec}
 import org.apache.arrow.memory.{ArrowBuf, BufferAllocator}
+import org.apache.arrow.util.AutoCloseables
+import org.apache.arrow.vector.{VectorSchemaRoot, VectorUnloader}
 import org.apache.arrow.vector.compression.{CompressionCodec, CompressionUtil, NoCompressionCodec}
+import org.apache.arrow.vector.ipc.message.ArrowRecordBatch
 
 import org.apache.spark.SparkException
 
@@ -52,6 +58,23 @@ private[sql] object ArrowCompressionUtils {
           s"Unsupported Arrow compression codec: $other. Supported values: none, zstd, lz4")
     }
   }
+
+  /**
+   * Creates the [[VectorUnloader]] that turns `root` into record batches compressed with `codec`,
+   * with null counts and aligned buffers.
+   *
+   * With a compressing codec, the unloader releases what it holds for a batch when compressing
+   * the batch fails, e.g. on a direct memory OOM. Arrow Java before 20.0.0 leaks the buffer being
+   * compressed and the buffers already compressed for the batch (apache/arrow-java#1234), and
+   * they stay allocated until the JVM exits.
+   */
+  def createVectorUnloader(root: VectorSchemaRoot, codec: CompressionCodec): VectorUnloader = {
+    if (codec.getCodecType == CompressionUtil.CodecType.NO_COMPRESSION) {
+      new VectorUnloader(root, true, codec, true)
+    } else {
+      new ReleasingVectorUnloader(root, new TrackingCompressionCodec(codec))
+    }
+  }
 }
 
 /**
@@ -83,3 +106,43 @@ private[sql] class BoundedZstdCompressionCodec(compressionLevel: Int)
     compressedBuffer
   }
 }
+
+/**
+ * A [[VectorUnloader]] that releases the buffers it holds for a batch it fails to unload. Remove
+ * once on Arrow 20.0.0, which releases them itself, so they would be released twice
+ * (apache/arrow-java#1290).
+ */
+private class ReleasingVectorUnloader(root: VectorSchemaRoot, codec: TrackingCompressionCodec)
+  extends VectorUnloader(root, true, codec, true) {
+
+  override def getRecordBatch(): ArrowRecordBatch = {
+    try {
+      super.getRecordBatch()
+    } catch {
+      case t: Throwable =>
+        AutoCloseables.close(t, codec.heldBuffers.asJava)
+        throw t
+    } finally {
+      // A returned batch owns the buffers now.
+      codec.heldBuffers.clear()
+    }
+  }
+}
+
+/** Delegates to `codec` and keeps, for each call, the buffer that the batch holds afterwards. */
+private class TrackingCompressionCodec(codec: CompressionCodec) extends CompressionCodec {
+  val heldBuffers = new ArrayBuffer[ArrowBuf]
+
+  override def compress(allocator: BufferAllocator, input: ArrowBuf): ArrowBuf = {
+    // The unloader retained `input` for `codec`, which releases it only when it succeeds.
+    heldBuffers += input
+    val output = codec.compress(allocator, input)
+    heldBuffers(heldBuffers.length - 1) = output
+    output
+  }
+
+  override def decompress(allocator: BufferAllocator, input: ArrowBuf): ArrowBuf =
+    codec.decompress(allocator, input)
+
+  override def getCodecType: CompressionUtil.CodecType = codec.getCodecType
+}

sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowConverters.scala

@@ -99,7 +99,7 @@ private[sql] object ArrowConverters extends Logging {
     // Create compression codec based on config
     private val codec = ArrowCompressionUtils.createCompressionCodec(
       SQLConf.get.arrowCompressionCodec, SQLConf.get.arrowZstdCompressionLevel)
-    protected val unloader = new VectorUnloader(root, true, codec, true)
+    protected val unloader = ArrowCompressionUtils.createVectorUnloader(root, codec)
     protected val arrowWriter = ArrowWriter.create(root)
     protected def maxBytesPerBatch: Long = -1L
 

sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala

@@ -24,7 +24,7 @@ import scala.jdk.CollectionConverters._
 
 import org.apache.arrow.flatbuf.{RecordBatch => FlatBufRecordBatch}
 import org.apache.arrow.memory.BufferAllocator
-import org.apache.arrow.vector.{TypeLayout, VectorLoader, VectorSchemaRoot, VectorUnloader}
+import org.apache.arrow.vector.{TypeLayout, VectorLoader, VectorSchemaRoot}
 import org.apache.arrow.vector.compression.NoCompressionCodec
 import org.apache.arrow.vector.ipc.{ReadChannel, WriteChannel}
 import org.apache.arrow.vector.ipc.message.{ArrowBodyCompression, ArrowFieldNode}
@@ -998,7 +998,7 @@ private class InternalRowToArrowCachedBatchIterator(
     sparkSchema, timeZoneId, false, false, losslessInternalTypes = true)
   private val root = VectorSchemaRoot.create(arrowSchema, allocator)
   private val arrowWriter = ArrowWriter.create(root)
-  private val unloader = new VectorUnloader(root, true, compressionCodec, true)
+  private val unloader = ArrowCompressionUtils.createVectorUnloader(root, compressionCodec)
 
   // Create statistics collectors for each column
   private val statsCollectors: Array[ColumnStats] = schema.map { attr =>
@@ -1177,7 +1177,7 @@ private class ColumnarBatchToArrowCachedBatchIterator(
 
     Utils.tryWithSafeFinally {
       // Use VectorUnloader to create compressed RecordBatch
-      val unloader = new VectorUnloader(root, true, compressionCodec, true)
+      val unloader = ArrowCompressionUtils.createVectorUnloader(root, compressionCodec)
       val recordBatch = unloader.getRecordBatch()
 
       Utils.tryWithSafeFinally {
@@ -1204,7 +1204,7 @@ private class ColumnarBatchToArrowCachedBatchIterator(
     // Convert columnar batch to rows, then to Arrow
     val root = VectorSchemaRoot.create(arrowSchema, allocator)
     val arrowWriter = ArrowWriter.create(root)
-    val unloader = new VectorUnloader(root, true, compressionCodec, true)
+    val unloader = ArrowCompressionUtils.createVectorUnloader(root, compressionCodec)
 
     Utils.tryWithSafeFinally {
       val rowIterator = batch.rowIterator().asScala

sql/core/src/main/scala/org/apache/spark/sql/execution/python/CoGroupedArrowPythonRunner.scala

@@ -94,7 +94,7 @@ class CoGroupedArrowPythonRunner(
   private def createUnloader(root: VectorSchemaRoot): VectorUnloader = {
     val codec = ArrowCompressionUtils.createCompressionCodec(
       compressionCodecName, SQLConf.get.arrowZstdCompressionLevel)
-    new VectorUnloader(root, true, codec, true)
+    ArrowCompressionUtils.createVectorUnloader(root, codec)
   }
 
   protected def newWriter(

sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonInput.scala

@@ -27,7 +27,7 @@ import org.apache.arrow.vector.ipc.WriteChannel
 import org.apache.arrow.vector.ipc.message.MessageSerializer
 
 import org.apache.spark.api.python.BasePythonRunner
-import org.apache.spark.sql.execution.arrow.ArrowWriter
+import org.apache.spark.sql.execution.arrow.{ArrowCompressionUtils, ArrowWriter}
 import org.apache.spark.sql.util.ArrowUtils
 import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch}
 
@@ -55,7 +55,8 @@ private[python] trait ColumnarArrowPythonInput
   // Lazy-initialized ArrowWriter for the fallback (non-Arrow) path.
   // Uses the pre-allocated `root` from PythonArrowInput.
   private lazy val arrowWriter: ArrowWriter = ArrowWriter.create(root)
-  private lazy val unloader: VectorUnloader = new VectorUnloader(root, true, codec, true)
+  private lazy val unloader: VectorUnloader =
+    ArrowCompressionUtils.createVectorUnloader(root, codec)
 
   protected def writeNextBatchToArrowStream(
       root: VectorSchemaRoot,
@@ -119,7 +120,7 @@ private[python] trait ColumnarArrowPythonInput
     val batchRoot = VectorSchemaRoot.of(selectedVectors: _*)
     batchRoot.setRowCount(batch.numRows())
 
-    val batchUnloader = new VectorUnloader(batchRoot, true, codec, true)
+    val batchUnloader = ArrowCompressionUtils.createVectorUnloader(batchRoot, codec)
     val recordBatch = batchUnloader.getRecordBatch()
     try {
       val writeChannel = new WriteChannel(Channels.newChannel(dataOut))

sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowInput.scala

@@ -135,7 +135,7 @@ private[python] trait PythonArrowInput[IN] { self: BasePythonRunner[IN, _] =>
 private[python] trait BasicPythonArrowInput extends PythonArrowInput[Iterator[InternalRow]] {
   self: BasePythonRunner[Iterator[InternalRow], _] =>
   protected lazy val arrowWriter: arrow.ArrowWriter = ArrowWriter.create(root)
-  protected lazy val unloader = new VectorUnloader(root, true, codec, true)
+  protected lazy val unloader = ArrowCompressionUtils.createVectorUnloader(root, codec)
 
   protected val maxRecordsPerBatch: Int = {
     val v = SQLConf.get.arrowMaxRecordsPerBatch
@@ -295,7 +295,7 @@ private[python] trait GroupedPythonArrowInput { self: RowInputArrowPythonRunner
               schema, timeZoneId, pythonExec,
               largeVarTypes, dataOut, context)
             // Set the unloader with compression after creating the writer
-            writer.unloader = new VectorUnloader(writer.root, true, self.codec, true)
+            writer.unloader = ArrowCompressionUtils.createVectorUnloader(writer.root, self.codec)
             nextBatchStart = inputIterator.next()
           }
         }

sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtilsSuite.scala

@@ -21,6 +21,7 @@ import java.io.ByteArrayOutputStream
 import java.nio.{ByteBuffer, ByteOrder}
 import java.nio.channels.Channels
 
+import scala.jdk.CollectionConverters._
 import scala.util.Random
 
 import org.apache.arrow.compression.ZstdCompressionCodec
@@ -260,4 +261,44 @@ class ArrowCompressionUtilsSuite extends SparkFunSuite {
       allocator.close()
     }
   }
+
+  test("SPARK-60776: unloader releases the buffers of a batch it fails to unload") {
+    var calls = 0
+    val codec = new BoundedZstdCompressionCodec(3) {
+      override protected def doCompress(allocator: BufferAllocator, input: ArrowBuf): ArrowBuf = {
+        calls += 1
+        if (calls == 4) throw new IllegalStateException("compression failed")
+        super.doCompress(allocator, input)
+      }
+    }
+    val arrowSchema = ArrowUtils.toArrowSchema(
+      StructType(Seq(StructField("c", LongType))), "UTC", true, false)
+    val allocator =
+      ArrowUtils.rootAllocator.newChildAllocator(this.getClass.getSimpleName, 0, Long.MaxValue)
+    val root = VectorSchemaRoot.create(arrowSchema, allocator)
+    val writer = ArrowWriter.create(root)
+    val unloader = ArrowCompressionUtils.createVectorUnloader(root, codec)
+    Utils.tryWithSafeFinally {
+      writer.write(InternalRow(1L))
+      writer.finish()
+      // A batch compresses a validity buffer and a data buffer, so the first batch succeeds and the
+      // second fails holding its compressed validity buffer and its retained data buffer.
+      val batch = unloader.getRecordBatch()
+      Utils.tryWithSafeFinally {
+        val e = intercept[IllegalStateException](unloader.getRecordBatch())
+        assert(e.getMessage === "compression failed")
+        // A suppressed exception means that a buffer was released twice. If this fails after an
+        // Arrow upgrade, Arrow releases these buffers itself and ReleasingVectorUnloader must be
+        // removed.
+        assert(e.getSuppressed.isEmpty)
+        assert(batch.getBuffers.asScala.forall(_.refCnt() == 1))
+      } {
+        batch.close()
+      }
+    } {
+      // Closing the allocator fails if a buffer is left behind.
+      root.close()
+      allocator.close()
+    }
+  }
 }