GH-44910: [Swift] fix ipc stream reader and writer impl

abandy · abandy · commit f9e029e92f08 · 2025-01-16T21:23:46.000-05:00
diff --git a/swift/Arrow/Sources/Arrow/ArrowReader.swift b/swift/Arrow/Sources/Arrow/ArrowReader.swift
@@ -19,7 +19,7 @@ import FlatBuffers
 import Foundation
 
 let FILEMARKER = "ARROW1"
-let CONTINUATIONMARKER = -1
+let CONTINUATIONMARKER = UInt32(0xFFFFFFFF)
 
 public class ArrowReader { // swiftlint:disable:this type_body_length
     private class RecordBatchData {
@@ -216,7 +216,77 @@ public class ArrowReader { // swiftlint:disable:this type_body_length
         return .success(RecordBatch(arrowSchema, columns: columns))
     }
 
-    public func fromStream( // swiftlint:disable:this function_body_length
+    /*
+     The Memory stream format is for reading the arrow streaming protocol.  This
+     format is slightly different from the File format protocol as it doesn't contain
+     a header and footer
+     */
+    public func fromMemoryStream( // swiftlint:disable:this function_body_length
+        _ fileData: Data,
+        useUnalignedBuffers: Bool = false
+    ) -> Result<ArrowReaderResult, ArrowError> {
+        let result = ArrowReaderResult()
+        var offset: Int = 0
+        var length = getUInt32(fileData, offset: offset)
+        var streamData = fileData
+        var schemaMessage: org_apache_arrow_flatbuf_Schema?
+        while length != 0 {
+            if length == CONTINUATIONMARKER {
+                offset += Int(MemoryLayout<Int32>.size)
+                length = getUInt32(fileData, offset: offset)
+                if length == 0 {
+                    return .success(result)
+                }
+            }
+
+            offset += Int(MemoryLayout<Int32>.size)
+            streamData = fileData[offset...]
+            let dataBuffer = ByteBuffer(
+                data: streamData,
+                allowReadingUnalignedBuffers: true)
+            let message = org_apache_arrow_flatbuf_Message.getRootAsMessage(bb: dataBuffer)
+            switch message.headerType {
+            case .recordbatch:
+                do {
+                    let rbMessage = message.header(type: org_apache_arrow_flatbuf_RecordBatch.self)!
+                    offset += Int(message.bodyLength + Int64(length))
+                    let recordBatch = try loadRecordBatch(
+                        rbMessage,
+                        schema: schemaMessage!,
+                        arrowSchema: result.schema!,
+                        data: fileData,
+                        messageEndOffset: (message.bodyLength + Int64(length))).get()
+                    result.batches.append(recordBatch)
+                    length = getUInt32(fileData, offset: offset)
+                } catch let error as ArrowError {
+                    return .failure(error)
+                } catch {
+                    return .failure(.unknownError("Unexpected error: \(error)"))
+                }
+            case .schema:
+                schemaMessage = message.header(type: org_apache_arrow_flatbuf_Schema.self)!
+                let schemaResult = loadSchema(schemaMessage!)
+                switch schemaResult {
+                case .success(let schema):
+                    result.schema = schema
+                case .failure(let error):
+                    return .failure(error)
+                }
+                offset += Int(message.bodyLength + Int64(length))
+                length = getUInt32(fileData, offset: offset)
+            default:
+                return .failure(.unknownError("Unhandled header type: \(message.headerType)"))
+            }
+        }
+        return .success(result)
+    }
+
+    /*
+     The File stream format supports random accessing the data.  This format contains
+     a header and footer around the streaming format.
+     */
+
+    public func fromFileStream( // swiftlint:disable:this function_body_length
         _ fileData: Data,
         useUnalignedBuffers: Bool = false
     ) -> Result<ArrowReaderResult, ArrowError> {
@@ -242,7 +312,7 @@ public class ArrowReader { // swiftlint:disable:this type_body_length
         for index in 0 ..< footer.recordBatchesCount {
             let recordBatch = footer.recordBatches(at: index)!
             var messageLength = fileData.withUnsafeBytes { rawBuffer in
-                rawBuffer.loadUnaligned(fromByteOffset: Int(recordBatch.offset), as: Int32.self)
+                rawBuffer.loadUnaligned(fromByteOffset: Int(recordBatch.offset), as: UInt32.self)
             }
 
             var messageOffset: Int64 = 1
@@ -251,7 +321,7 @@ public class ArrowReader { // swiftlint:disable:this type_body_length
                 messageLength = fileData.withUnsafeBytes { rawBuffer in
                     rawBuffer.loadUnaligned(
                         fromByteOffset: Int(recordBatch.offset + Int64(MemoryLayout<Int32>.size)),
-                        as: Int32.self)
+                        as: UInt32.self)
                 }
             }
 
@@ -296,7 +366,7 @@ public class ArrowReader { // swiftlint:disable:this type_body_length
             let markerLength = FILEMARKER.utf8.count
             let footerLengthEnd = Int(fileData.count - markerLength)
             let data = fileData[..<(footerLengthEnd)]
-            return fromStream(data)
+            return fromFileStream(data)
         } catch {
             return .failure(.unknownError("Error loading file: \(error)"))
         }
@@ -340,10 +410,10 @@ public class ArrowReader { // swiftlint:disable:this type_body_length
             } catch {
                 return .failure(.unknownError("Unexpected error: \(error)"))
             }
-
         default:
             return .failure(.unknownError("Unhandled header type: \(message.headerType)"))
         }
     }
 
 }
+// swiftlint:disable:this file_length
diff --git a/swift/Arrow/Sources/Arrow/ArrowReaderHelper.swift b/swift/Arrow/Sources/Arrow/ArrowReaderHelper.swift
@@ -289,3 +289,10 @@ func validateFileData(_ data: Data) -> Bool {
     let endString = String(decoding: data[(data.count - markerLength)...], as: UTF8.self)
     return startString == FILEMARKER && endString == FILEMARKER
 }
+
+func getUInt32(_ data: Data, offset: Int) -> UInt32 {
+    let token = data.withUnsafeBytes { rawBuffer in
+        rawBuffer.loadUnaligned(fromByteOffset: offset, as: UInt32.self)
+    }
+    return token
+}
diff --git a/swift/Arrow/Sources/Arrow/ArrowWriter.swift b/swift/Arrow/Sources/Arrow/ArrowWriter.swift
@@ -123,6 +123,7 @@ public class ArrowWriter { // swiftlint:disable:this type_body_length
             let startIndex = writer.count
             switch writeRecordBatch(batch: batch) {
             case .success(let rbResult):
+                withUnsafeBytes(of: CONTINUATIONMARKER.littleEndian) {writer.append(Data($0))}
                 withUnsafeBytes(of: rbResult.1.o.littleEndian) {writer.append(Data($0))}
                 writer.append(rbResult.0)
                 switch writeRecordBatchData(&writer, batch: batch) {
@@ -232,7 +233,7 @@ public class ArrowWriter { // swiftlint:disable:this type_body_length
         return .success(fbb.data)
     }
 
-    private func writeStream(_ writer: inout DataWriter, info: ArrowWriter.Info) -> Result<Bool, ArrowError> {
+    private func writeFileStream(_ writer: inout DataWriter, info: ArrowWriter.Info) -> Result<Bool, ArrowError> {
         var fbb: FlatBufferBuilder = FlatBufferBuilder()
         switch writeSchema(&fbb, schema: info.schema) {
         case .success(let schemaOffset):
@@ -264,9 +265,41 @@ public class ArrowWriter { // swiftlint:disable:this type_body_length
         return .success(true)
     }
 
-    public func toStream(_ info: ArrowWriter.Info) -> Result<Data, ArrowError> {
+    public func toMemoryStream(_ info: ArrowWriter.Info) -> Result<Data, ArrowError> {
+        let writer: any DataWriter = InMemDataWriter()
+        switch toMessage(info.schema) {
+        case .success(let schemaData):
+            withUnsafeBytes(of: CONTINUATIONMARKER.littleEndian) {writer.append(Data($0))}
+            withUnsafeBytes(of: UInt32(schemaData.count).littleEndian) {writer.append(Data($0))}
+            writer.append(schemaData)
+        case .failure(let error):
+            return .failure(error)
+        }
+
+        for batch in info.batches {
+            switch toMessage(batch) {
+            case .success(let batchData):
+                withUnsafeBytes(of: CONTINUATIONMARKER.littleEndian) {writer.append(Data($0))}
+                withUnsafeBytes(of: UInt32(batchData[0].count).littleEndian) {writer.append(Data($0))}
+                writer.append(batchData[0])
+                writer.append(batchData[1])
+            case .failure(let error):
+                return .failure(error)
+            }
+        }
+
+        withUnsafeBytes(of: CONTINUATIONMARKER.littleEndian) {writer.append(Data($0))}
+        withUnsafeBytes(of: UInt32(0).littleEndian) {writer.append(Data($0))}
+        if let memWriter = writer as? InMemDataWriter {
+            return .success(memWriter.data)
+        } else {
+            return .failure(.invalid("Unable to cast writer"))
+        }
+    }
+
+    public func toFileStream(_ info: ArrowWriter.Info) -> Result<Data, ArrowError> {
         var writer: any DataWriter = InMemDataWriter()
-        switch writeStream(&writer, info: info) {
+        switch writeFileStream(&writer, info: info) {
         case .success:
             if let memWriter = writer as? InMemDataWriter {
                 return .success(memWriter.data)
@@ -293,7 +326,7 @@ public class ArrowWriter { // swiftlint:disable:this type_body_length
 
         var writer: any DataWriter = FileDataWriter(fileHandle)
         writer.append(FILEMARKER.data(using: .utf8)!)
-        switch writeStream(&writer, info: info) {
+        switch writeFileStream(&writer, info: info) {
         case .success:
             writer.append(FILEMARKER.data(using: .utf8)!)
         case .failure(let error):
diff --git a/swift/Arrow/Tests/ArrowTests/IPCTests.swift b/swift/Arrow/Tests/ArrowTests/IPCTests.swift
@@ -167,10 +167,10 @@ final class IPCFileReaderTests: XCTestCase {
         let arrowWriter = ArrowWriter()
         // write data from file to a stream
         let writerInfo = ArrowWriter.Info(.recordbatch, schema: fileRBs[0].schema, batches: fileRBs)
-        switch arrowWriter.toStream(writerInfo) {
+        switch arrowWriter.toFileStream(writerInfo) {
         case .success(let writeData):
             // read stream back into recordbatches
-            try checkBoolRecordBatch(arrowReader.fromStream(writeData))
+            try checkBoolRecordBatch(arrowReader.fromFileStream(writeData))
         case .failure(let error):
             throw error
         }
@@ -190,10 +190,10 @@ final class IPCFileReaderTests: XCTestCase {
         let recordBatch = try makeRecordBatch()
         let arrowWriter = ArrowWriter()
         let writerInfo = ArrowWriter.Info(.recordbatch, schema: schema, batches: [recordBatch])
-        switch arrowWriter.toStream(writerInfo) {
+        switch arrowWriter.toFileStream(writerInfo) {
         case .success(let writeData):
             let arrowReader = ArrowReader()
-            switch arrowReader.fromStream(writeData) {
+            switch arrowReader.fromFileStream(writeData) {
             case .success(let result):
                 let recordBatches = result.batches
                 XCTAssertEqual(recordBatches.count, 1)
@@ -242,10 +242,10 @@ final class IPCFileReaderTests: XCTestCase {
         let schema = makeSchema()
         let arrowWriter = ArrowWriter()
         let writerInfo = ArrowWriter.Info(.schema, schema: schema)
-        switch arrowWriter.toStream(writerInfo) {
+        switch arrowWriter.toFileStream(writerInfo) {
         case .success(let writeData):
             let arrowReader = ArrowReader()
-            switch arrowReader.fromStream(writeData) {
+            switch arrowReader.fromFileStream(writeData) {
             case .success(let result):
                 XCTAssertNotNil(result.schema)
                 let schema  = result.schema!
@@ -325,10 +325,10 @@ final class IPCFileReaderTests: XCTestCase {
         let dataset = try makeBinaryDataset()
         let writerInfo = ArrowWriter.Info(.recordbatch, schema: dataset.0, batches: [dataset.1])
         let arrowWriter = ArrowWriter()
-        switch arrowWriter.toStream(writerInfo) {
+        switch arrowWriter.toFileStream(writerInfo) {
         case .success(let writeData):
             let arrowReader = ArrowReader()
-            switch arrowReader.fromStream(writeData) {
+            switch arrowReader.fromFileStream(writeData) {
             case .success(let result):
                 XCTAssertNotNil(result.schema)
                 let schema  = result.schema!
@@ -354,10 +354,10 @@ final class IPCFileReaderTests: XCTestCase {
         let dataset = try makeTimeDataset()
         let writerInfo = ArrowWriter.Info(.recordbatch, schema: dataset.0, batches: [dataset.1])
         let arrowWriter = ArrowWriter()
-        switch arrowWriter.toStream(writerInfo) {
+        switch arrowWriter.toFileStream(writerInfo) {
         case .success(let writeData):
             let arrowReader = ArrowReader()
-            switch arrowReader.fromStream(writeData) {
+            switch arrowReader.fromFileStream(writeData) {
             case .success(let result):
                 XCTAssertNotNil(result.schema)
                 let schema  = result.schema!
