This is an automated email from the ASF dual-hosted git repository. raulcd pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/arrow-swift.git
commit ebe3512d850012e19d5e2f52559a590b575ee584 Author: abandy <[email protected]> AuthorDate: Fri Dec 8 14:50:22 2023 -0700 GH-37884: [Swift] allow reading of unaligned FlatBuffers buffers (#38635) The PR enables the swift readers to read from unaligned buffers (fix for issue: 37884) Enabling unaligned buffers incurs a performance penalty so the developer will need to consider this when enabling this feature. It is not currently possible to recover from a buffer unaligned error as this error is a fatalError so trying aligned and then falling back to unaligned is not an option. Also, FlatBuffers has a verifier that should be able to catch this error but currently it seems to fail on both aligned and unaligned buffers (I tried verifying the example python server get return value but verification fails even though the buffers are able to be read successfully) * Closes: #37884 Authored-by: Alva Bandy <[email protected]> Signed-off-by: Sutou Kouhei <[email protected]> --- Arrow/Package.swift | 6 +++++- Arrow/Sources/Arrow/ArrowReader.swift | 18 +++++++++++++----- ArrowFlight/Sources/ArrowFlight/FlightClient.swift | 11 +++++++++-- ArrowFlight/Sources/ArrowFlight/FlightServer.swift | 7 +++++++ .../Sources/ArrowFlight/RecordBatchStreamReader.swift | 11 +++++++++-- 5 files changed, 43 insertions(+), 10 deletions(-) diff --git a/Arrow/Package.swift b/Arrow/Package.swift index 065afe6..946eb99 100644 --- a/Arrow/Package.swift +++ b/Arrow/Package.swift @@ -32,7 +32,11 @@ let package = Package( targets: ["Arrow"]), ], dependencies: [ - .package(url: "https://github.com/google/flatbuffers.git";, from: "23.3.3") + // The latest version of flatbuffers v23.5.26 was built in May 26, 2023 + // and therefore doesn't include the unaligned buffer swift changes. + // This can be changed back to using the tag once a new version of + // flatbuffers has been released. + .package(url: "https://github.com/google/flatbuffers.git";, branch: "master") ], targets: [ // Targets are the basic building blocks of a package. A target can define a module or a test suite. diff --git a/Arrow/Sources/Arrow/ArrowReader.swift b/Arrow/Sources/Arrow/ArrowReader.swift index ef995b1..d9dc1bd 100644 --- a/Arrow/Sources/Arrow/ArrowReader.swift +++ b/Arrow/Sources/Arrow/ArrowReader.swift @@ -132,7 +132,8 @@ public class ArrowReader { } public func fromStream( // swiftlint:disable:this function_body_length - _ fileData: Data + _ fileData: Data, + useUnalignedBuffers: Bool = false ) -> Result<ArrowReaderResult, ArrowError> { let footerLength = fileData.withUnsafeBytes { rawBuffer in rawBuffer.loadUnaligned(fromByteOffset: fileData.count - 4, as: Int32.self) @@ -141,7 +142,9 @@ public class ArrowReader { let result = ArrowReaderResult() let footerStartOffset = fileData.count - Int(footerLength + 4) let footerData = fileData[footerStartOffset...] - let footerBuffer = ByteBuffer(data: footerData) + let footerBuffer = ByteBuffer( + data: footerData, + allowReadingUnalignedBuffers: useUnalignedBuffers) let footer = org_apache_arrow_flatbuf_Footer.getRootAsFooter(bb: footerBuffer) let schemaResult = loadSchema(footer.schema!) switch schemaResult { @@ -170,7 +173,9 @@ public class ArrowReader { let messageStartOffset = recordBatch.offset + (Int64(MemoryLayout<Int32>.size) * messageOffset) let messageEndOffset = messageStartOffset + Int64(messageLength) let recordBatchData = fileData[messageStartOffset ..< messageEndOffset] - let mbb = ByteBuffer(data: recordBatchData) + let mbb = ByteBuffer( + data: recordBatchData, + allowReadingUnalignedBuffers: useUnalignedBuffers) let message = org_apache_arrow_flatbuf_Message.getRootAsMessage(bb: mbb) switch message.headerType { case .recordbatch: @@ -219,9 +224,12 @@ public class ArrowReader { public func fromMessage( _ dataHeader: Data, dataBody: Data, - result: ArrowReaderResult + result: ArrowReaderResult, + useUnalignedBuffers: Bool = false ) -> Result<Void, ArrowError> { - let mbb = ByteBuffer(data: dataHeader) + let mbb = ByteBuffer( + data: dataHeader, + allowReadingUnalignedBuffers: useUnalignedBuffers) let message = org_apache_arrow_flatbuf_Message.getRootAsMessage(bb: mbb) switch message.headerType { case .schema: diff --git a/ArrowFlight/Sources/ArrowFlight/FlightClient.swift b/ArrowFlight/Sources/ArrowFlight/FlightClient.swift index 7a572ce..ef3e4fa 100644 --- a/ArrowFlight/Sources/ArrowFlight/FlightClient.swift +++ b/ArrowFlight/Sources/ArrowFlight/FlightClient.swift @@ -24,8 +24,11 @@ import Arrow public class FlightClient { let client: Arrow_Flight_Protocol_FlightServiceAsyncClient - public init(channel: GRPCChannel) { + let allowReadingUnalignedBuffers: Bool + + public init(channel: GRPCChannel, allowReadingUnalignedBuffers: Bool = false ) { client = Arrow_Flight_Protocol_FlightServiceAsyncClient(channel: channel) + self.allowReadingUnalignedBuffers = allowReadingUnalignedBuffers } private func readMessages( @@ -34,7 +37,11 @@ public class FlightClient { let reader = ArrowReader() let arrowResult = ArrowReader.makeArrowReaderResult() for try await data in responseStream { - switch reader.fromMessage(data.dataHeader, dataBody: data.dataBody, result: arrowResult) { + switch reader.fromMessage( + data.dataHeader, + dataBody: data.dataBody, + result: arrowResult, + useUnalignedBuffers: allowReadingUnalignedBuffers) { case .success: continue case .failure(let error): diff --git a/ArrowFlight/Sources/ArrowFlight/FlightServer.swift b/ArrowFlight/Sources/ArrowFlight/FlightServer.swift index a34bf5c..19644d6 100644 --- a/ArrowFlight/Sources/ArrowFlight/FlightServer.swift +++ b/ArrowFlight/Sources/ArrowFlight/FlightServer.swift @@ -63,6 +63,7 @@ public func schemaFromMessage(_ schemaData: Data) -> ArrowSchema? { } public protocol ArrowFlightServer: Sendable { + var allowReadingUnalignedBuffers: Bool { get } func listFlights(_ criteria: FlightCriteria, writer: FlightInfoStreamWriter) async throws func getFlightInfo(_ request: FlightDescriptor) async throws -> FlightInfo func getSchema(_ request: FlightDescriptor) async throws -> ArrowFlight.FlightSchemaResult @@ -73,6 +74,12 @@ public protocol ArrowFlightServer: Sendable { func doExchange(_ reader: RecordBatchStreamReader, writer: RecordBatchStreamWriter) async throws } +extension ArrowFlightServer { + var allowReadingUnalignedBuffers: Bool { + return false + } +} + public func makeFlightServer(_ handler: ArrowFlightServer) -> CallHandlerProvider { return InternalFlightServer(handler) } diff --git a/ArrowFlight/Sources/ArrowFlight/RecordBatchStreamReader.swift b/ArrowFlight/Sources/ArrowFlight/RecordBatchStreamReader.swift index 972d194..464752d 100644 --- a/ArrowFlight/Sources/ArrowFlight/RecordBatchStreamReader.swift +++ b/ArrowFlight/Sources/ArrowFlight/RecordBatchStreamReader.swift @@ -27,10 +27,13 @@ public class RecordBatchStreamReader: AsyncSequence, AsyncIteratorProtocol { var descriptor: FlightDescriptor? var batchIndex = 0 var streamIterator: any AsyncIteratorProtocol + var useUnalignedBuffers: Bool let stream: GRPC.GRPCAsyncRequestStream<Arrow_Flight_Protocol_FlightData> - init(_ stream: GRPC.GRPCAsyncRequestStream<Arrow_Flight_Protocol_FlightData>) { + init(_ stream: GRPC.GRPCAsyncRequestStream<Arrow_Flight_Protocol_FlightData>, + useUnalignedBuffers: Bool = false) { self.stream = stream self.streamIterator = self.stream.makeAsyncIterator() + self.useUnalignedBuffers = useUnalignedBuffers } public func next() async throws -> (Arrow.RecordBatch?, FlightDescriptor?)? { @@ -55,7 +58,11 @@ public class RecordBatchStreamReader: AsyncSequence, AsyncIteratorProtocol { let dataBody = flightData.dataBody let dataHeader = flightData.dataHeader descriptor = FlightDescriptor(flightData.flightDescriptor) - switch reader.fromMessage(dataHeader, dataBody: dataBody, result: result) { + switch reader.fromMessage( + dataHeader, + dataBody: dataBody, + result: result, + useUnalignedBuffers: useUnalignedBuffers) { case .success(()): if result.batches.count > 0 { batches = result.batches
