(spark-connect-swift) branch main updated: [SPARK-60472] Fix `collect` to throw an error instead of misreading compressed Arrow batches

[do...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

dongjoon-hyun pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/spark-connect-swift.git


The following commit(s) were added to refs/heads/main by this push:
     new 9b32afb  [SPARK-60472] Fix `collect` to throw an error instead of misreading compressed Arrow batches
9b32afb is described below

commit 9b32afb97425282be5c796625bc9af1129967016
Author: Dongjoon Hyun <do...@apache.org>
AuthorDate: Thu Oct 22 08:34:35 2026 -0700

    [SPARK-60472] Fix `collect` to throw an error instead of misreading compressed Arrow batches
    
    ### What changes were proposed in this pull request?
    
    This PR fixes `DataFrame` actions to throw `SparkConnectError.unsupportedOperation` instead of returning wrong values or crashing when the Spark Connect server sends compressed Arrow batches.
    
    - `ArrowReader` returns an error for a record batch with `compression` instead of reading the compressed buffers as they are.
    - `DataFrame` actions check the header of each Arrow batch first and throw `unsupportedOperation` with the codec name and the workaround, `spark.sql.execution.arrow.compression.codec=none`.
    
    ### Why are the changes needed?
    
    Since Spark 4.1.0, the server compresses the buffers of Arrow batches when [spark.sql.execution.arrow.compression.codec](https://github.com/apache/spark/blob/v4.1.0/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala#L4046) is `zstd` or `lz4`. [ArrowBatchIterator](https://github.com/apache/spark/blob/v4.1.0/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowConverters.scala#L118) of `v4.1.0` (2025-12-11) passes the codec to `VectorUnloader`. The serv [...]
    
    ```swift
    try await spark.conf.set("spark.sql.execution.arrow.compression.codec", "zstd")
    try await spark.range(10).count()    // -1
    try await spark.range(10).collect()  // Wrong values like -1
    try await spark.range(3).show()      // Fatal error: UnsafeBufferPointer with negative count
    ```
    
    Like the Scala client, which fails with `IllegalArgumentException: Please add arrow-compression module to use CommonsCompressionFactory for ZSTD`, this PR does not support decompression. PySpark decompresses the batches with PyArrow.
    
    ### Does this PR introduce _any_ user-facing change?
    
    Yes. With Spark 4.1+ and a compression codec, `collect()`, `collect(as:)`, `head()`, `first()`, `take()`, `tail()`, `count()`, and `show()` throw the following error instead of returning wrong values or crashing. Results without buffers, e.g., `df.select()`, also throw it.
    
    ```
    unsupportedOperation: Unsupported Arrow compression codec: zstd. Set `spark.sql.execution.arrow.compression.codec` to `none`.
    ```
    
    There is no change when the codec is `none`, which is the default.
    
    ### How was this patch tested?
    
    Added `RecordBatchTests.readCompressed`, which does not require a Spark Connect server, and `DataFrameTests.arrowCompression` for Spark 4.1+. Verified against an `apache/spark:4.2.0` Connect server.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Generated-by: Claude Opus 5.5
    
    Closes #618 from dongjoon-hyun/SPARK-60472.
    
    Authored-by: Dongjoon Hyun <do...@apache.org>
    Signed-off-by: Dongjoon Hyun <do...@apache.org>
---
 Sources/SparkConnect/ArrowReader.swift         |  4 ++++
 Sources/SparkConnect/DataFrame+Actions.swift   | 15 +++++++++++++++
 Tests/SparkConnectTests/DataFrameTests.swift   | 19 +++++++++++++++++++
 Tests/SparkConnectTests/RecordBatchTests.swift | 25 +++++++++++++++++++++++++
 4 files changed, 63 insertions(+)

diff --git a/Sources/SparkConnect/ArrowReader.swift b/Sources/SparkConnect/ArrowReader.swift
index cde7c28..4d31401 100644
--- a/Sources/SparkConnect/ArrowReader.swift
+++ b/Sources/SparkConnect/ArrowReader.swift
@@ -284,6 +284,10 @@ public class ArrowReader {  // swiftlint:disable:this type_body_length
     guard let length = UInt(exactly: recordBatch.length) else {
       return .failure(.invalid("Invalid record batch length: \(recordBatch.length)"))
     }
+    // Compressed buffers cannot be read as they are, and decompression is not supported.
+    if let compression = recordBatch.compression {
+      return .failure(.invalid("Unsupported record batch compression codec: \(compression.codec)"))
+    }
     var columns: [ArrowArrayHolder] = []
     let batchData = RecordBatchData(recordBatch, schema: schema)
     let loadInfo = DataLoadInfo(
diff --git a/Sources/SparkConnect/DataFrame+Actions.swift b/Sources/SparkConnect/DataFrame+Actions.swift
index 3dd484c..f6dd55a 100644
--- a/Sources/SparkConnect/DataFrame+Actions.swift
+++ b/Sources/SparkConnect/DataFrame+Actions.swift
@@ -22,6 +22,7 @@ import FoundationEssentials
 #else
 import Foundation
 #endif
+import FlatBuffers
 import GRPCProtobuf
 import Synchronization
 
@@ -100,6 +101,20 @@ extension DataFrame {
         let dataBodySize = totalSize - pos - 8
         let dataBody = Data(ipcStreamBytes[pos..<(pos + dataBodySize)])
 
+        // Spark 4.1+ compresses buffers if `spark.sql.execution.arrow.compression.codec` is set.
+        let message = org_apache_arrow_flatbuf_Message.getRootAsMessage(
+          bb: ByteBuffer(data: dataHeader))
+        if message.headerType == .recordbatch,
+          let recordBatch = message.header(type: org_apache_arrow_flatbuf_RecordBatch.self),
+          let compression = recordBatch.compression
+        {
+          throw SparkConnectError.unsupportedOperation(
+            SparkConnectError.Details(
+              message:
+                "Unsupported Arrow compression codec: \(compression.codec). "
+                + "Set `spark.sql.execution.arrow.compression.codec` to `none`."))
+        }
+
         // Read ArrowBatches
         try batches.withLock { batches in
           let reader = ArrowReader()
diff --git a/Tests/SparkConnectTests/DataFrameTests.swift b/Tests/SparkConnectTests/DataFrameTests.swift
index 8687a1b..7710d7a 100644
--- a/Tests/SparkConnectTests/DataFrameTests.swift
+++ b/Tests/SparkConnectTests/DataFrameTests.swift
@@ -774,6 +774,25 @@ struct DataFrameTests {
     await spark.stop()
   }
 
+  @Test
+  func arrowCompression() async throws {
+    let spark = try await SparkSession.builder.getOrCreate()
+    if await isSparkVersionAtLeast(spark.version, "4.1") {
+      let df = try await spark.range(10)
+      let expected = try await df.collect()
+      for codec in ["zstd", "lz4"] {
+        try await spark.conf.set("spark.sql.execution.arrow.compression.codec", codec)
+        let error = await #expect(throws: SparkConnectError.self) { try await df.collect() }
+        #expect(error == SparkConnectError.UnsupportedOperation)
+        #expect(error?.message.contains(codec) == true)
+        #expect(error?.message.contains("spark.sql.execution.arrow.compression.codec") == true)
+      }
+      try await spark.conf.unset("spark.sql.execution.arrow.compression.codec")
+      #expect(try await df.collect() == expected)
+    }
+    await spark.stop()
+  }
+
   @Test
   func isEmpty() async throws {
     let spark = try await SparkSession.builder.getOrCreate()
diff --git a/Tests/SparkConnectTests/RecordBatchTests.swift b/Tests/SparkConnectTests/RecordBatchTests.swift
index 639ec8d..d9fba66 100644
--- a/Tests/SparkConnectTests/RecordBatchTests.swift
+++ b/Tests/SparkConnectTests/RecordBatchTests.swift
@@ -22,6 +22,7 @@ import FoundationEssentials
 #else
 import Foundation
 #endif
+import FlatBuffers
 import Testing
 
 @testable import SparkConnect
@@ -54,4 +55,28 @@ struct RecordBatchTests {
     #expect(result.batches[0].columnCount == 0)
     #expect(result.batches[0].length == 3)
   }
+
+  @Test
+  func readCompressed() throws {
+    // Spark 4.1+ compresses the buffers if `spark.sql.execution.arrow.compression.codec` is set.
+    let schemaMessage = try ArrowWriter().toMessage(ArrowSchema.Builder().finish()).get()
+    for codec in [org_apache_arrow_flatbuf_CompressionType.lz4Frame, .zstd] {
+      var fbb = FlatBufferBuilder()
+      let compression = org_apache_arrow_flatbuf_BodyCompression.createBodyCompression(
+        &fbb, codec: codec)
+      let recordBatch = org_apache_arrow_flatbuf_RecordBatch.createRecordBatch(
+        &fbb, length: 3, compressionOffset: compression)
+      let message = org_apache_arrow_flatbuf_Message.createMessage(
+        &fbb, version: .max, headerType: .recordbatch, headerOffset: recordBatch)
+      fbb.finish(offset: message)
+
+      let reader = ArrowReader()
+      let result = ArrowReader.makeArrowReaderResult()
+      try reader.fromMessage(schemaMessage, dataBody: Data(), result: result).get()
+      #expect(throws: ArrowError.self) {
+        try reader.fromMessage(fbb.data, dataBody: Data(), result: result).get()
+      }
+      #expect(result.batches.isEmpty)
+    }
+  }
 }


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org