From 7fd1cea7ee1e84df03ba5ff29b169d279bfe047e Mon Sep 17 00:00:00 2001 From: Dongjoon Hyun Date: Mon, 2 Nov 2026 11:39:17 -0800 Subject: [PATCH] [SPARK-60816][SQL][PYTHON] Write empty Arrow buffers uncompressed so that PyArrow can read LZ4 batches --- python/pyspark/sql/tests/arrow/test_arrow.py | 9 ++++ .../arrow/ArrowCompressionUtils.scala | 30 +++++++++++- .../arrow/ArrowCompressionUtilsSuite.scala | 49 ++++++++++++++++++- 3 files changed, 84 insertions(+), 4 deletions(-) diff --git a/python/pyspark/sql/tests/arrow/test_arrow.py b/python/pyspark/sql/tests/arrow/test_arrow.py index df4239651d344..006bd9169d0b1 100644 --- a/python/pyspark/sql/tests/arrow/test_arrow.py +++ b/python/pyspark/sql/tests/arrow/test_arrow.py @@ -1894,6 +1894,15 @@ def test_toArrow_with_compression_codec_large_dataset(self): self.assertEqual(t.num_rows, 10000) self.assertEqual(t.column_names, ["id", "str_col", "mod_col"]) + def test_toArrow_with_compression_codec_empty_data_buffers(self): + # SPARK-60816: All strings are null, so the data buffer of the column is empty. + df = self.spark.range(3).selectExpr("CAST(NULL AS STRING) AS s") + + for codec in ["none", "zstd", "lz4"]: + with self.subTest(compressionCodec=codec): + with self.sql_conf({"spark.sql.execution.arrow.compression.codec": codec}): + self.assertEqual(df.toArrow().to_pydict(), {"s": [None] * 3}) + def test_toPandas_double_nested_array_empty_outer(self): schema = StructType([StructField("data", ArrayType(ArrayType(StringType())))]) df = self.spark.createDataFrame([Row(data=[])], schema=schema) 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 ce995bb0bc0d4..8c87dfc117b63 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 @@ -51,8 +51,9 @@ private[sql] object ArrowCompressionUtils { zstdCompressionLevel: Int): CompressionCodec = { codecName match { case "none" => NoCompressionCodec.INSTANCE - case "zstd" => new BoundedZstdCompressionCodec(zstdCompressionLevel) - case "lz4" => new Lz4CompressionCodec() + case "zstd" => + new BoundedZstdCompressionCodec(zstdCompressionLevel) with UncompressedEmptyBuffers + case "lz4" => new Lz4CompressionCodec with UncompressedEmptyBuffers case other => throw SparkException.internalError( s"Unsupported Arrow compression codec: $other. Supported values: none, zstd, lz4") @@ -146,6 +147,31 @@ private object BoundedZstdCompressionCodec { val MIN_SHRINK_SAVING_BYTES = 4 * 1024 } +/** + * A [[CompressionCodec]] mixin that gives an empty buffer the `-1` (stored uncompressed) length + * prefix. arrow-java up to 19.0 writes a `0` prefix and no compressed data, so Arrow C++ (and + * PyArrow) decompresses a zero-byte body, which LZ4 rejects with "Lz4 compressed input contains + * less than one frame". Remove once on Arrow 20.0.0 (apache/arrow-java#1248). + */ +private[sql] trait UncompressedEmptyBuffers extends CompressionCodec { + + abstract override def compress( + allocator: BufferAllocator, + uncompressedBuffer: ArrowBuf): ArrowBuf = { + if (uncompressedBuffer.writerIndex() == 0L) { + val compressedBuffer = allocator.buffer(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH) + compressedBuffer.setLong(0, CompressionUtil.NO_COMPRESSION_LENGTH) + compressedBuffer.writerIndex(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH) + // Like `super.compress`, release the input only on success. TrackingCompressionCodec relies + // on this to release it on failure. + uncompressedBuffer.close() + compressedBuffer + } else { + super.compress(allocator, uncompressedBuffer) + } + } +} + /** * 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 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 2bbdf770fea26..7737ccc350538 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 @@ -25,9 +25,9 @@ import scala.jdk.CollectionConverters._ import scala.util.Random import org.apache.arrow.compression.ZstdCompressionCodec -import org.apache.arrow.memory.{ArrowBuf, BufferAllocator, RootAllocator} +import org.apache.arrow.memory.{ArrowBuf, BufferAllocator, OutOfMemoryException, RootAllocator} import org.apache.arrow.vector.{VarCharVector, VectorSchemaRoot, VectorUnloader} -import org.apache.arrow.vector.compression.NoCompressionCodec +import org.apache.arrow.vector.compression.{CompressionUtil, NoCompressionCodec} import org.apache.arrow.vector.ipc.WriteChannel import org.apache.arrow.vector.ipc.message.MessageSerializer @@ -370,4 +370,49 @@ class ArrowCompressionUtilsSuite extends SparkFunSuite { assert(output.capacity() === allocatedSize) } } + + test("SPARK-60816: empty buffers are written with the uncompressed length prefix") { + // For an empty buffer, arrow-java up to 19.0 writes a `0` length prefix and no compressed data, + // which Arrow C++ (PyArrow) fails to decompress for LZ4. Readers already see the `-1` prefix + // for buffers that do not compress. + val allocator = + ArrowUtils.rootAllocator.newChildAllocator(this.getClass.getSimpleName, 0, Long.MaxValue) + Utils.tryWithSafeFinally { + Seq("lz4", "zstd").foreach { codecName => + val codec = ArrowCompressionUtils.createCompressionCodec(codecName, 3) + // A newly allocated buffer has writerIndex 0, so `compress` sees an empty buffer. + val compressed = codec.compress(allocator, allocator.buffer(16)) + try { + assert(compressed.writerIndex() === CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH) + assert(compressed.getLong(0) === CompressionUtil.NO_COMPRESSION_LENGTH, codecName) + } finally { + compressed.close() + } + } + } { + // Fails if `compress` leaked its input buffer. + allocator.close() + } + } + + test("SPARK-60816: unloader releases an empty buffer once if its prefix allocation fails") { + 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 codec = ArrowCompressionUtils.createCompressionCodec("lz4", 3) + val unloader = ArrowCompressionUtils.createVectorUnloader(root, codec) + Utils.tryWithSafeFinally { + root.allocateNew() + // With no rows the first buffer is empty, so its length prefix is the first allocation. + allocator.setLimit(0) + intercept[OutOfMemoryException](unloader.getRecordBatch()) + } { + // Closing the vector fails if the buffer was released twice, and closing the allocator fails + // if it was not released. + root.close() + allocator.close() + } + } }