diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtils.scala index e71680751ab8c..c6a57430ae48c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtils.scala +++ b/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 +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowConverters.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowConverters.scala index a593f4962c62c..e24048b3612d8 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowConverters.scala +++ b/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 diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala index db91f91e2480b..3a1e4829b2bef 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala +++ b/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 diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/CoGroupedArrowPythonRunner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/CoGroupedArrowPythonRunner.scala index 0a917677a52e6..1b3f0260de4e9 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/CoGroupedArrowPythonRunner.scala +++ b/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( diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonInput.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonInput.scala index fbc21c0e2bf6e..bee978f07fcf5 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonInput.scala +++ b/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)) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowInput.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowInput.scala index 596a330e2026c..6bc65a886abed 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowInput.scala +++ b/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() } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtilsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtilsSuite.scala index cde828e1c7437..f424aff29720f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowCompressionUtilsSuite.scala +++ b/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() + } + } }