[SPARK-60816][SQL][PYTHON] Write empty Arrow buffers uncompressed so that PyArrow can read LZ4 batches

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)

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

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