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 49681a139a0b7ef92b22e3a8e605a4677e10333b Author: abandy <[email protected]> AuthorDate: Thu Jun 27 22:34:28 2024 -0400 GH-42021: [Swift] Add Arrow encoder implementation for Swift Codable (#43063) ### Rationale for this change This change implements encoder for the Arrow Swift Codable. This will allow data in a struct/class to be copied to a RecordBatch. ### What changes are included in this PR? The ArrowEncoder class is included in this PR along with a class for each container type (keyed, unkeyed, and single). Most of the logic is encapsulated in the ArrowEncoder with minimal logic in each container class (Most of the methods in the container classes are a single line that calls methods in ArrowEncoder). - The ArrowEncoder does not allow nested encoding. The inherited methods for creating a nested encoder do not throw errors so an "errorMsg" member has been added to the ArrowEncoder to indicate that an error occurred in a method that does not throw. If this member is set, then an error will be thrown at the next available opportunity. - The underlaying Codable implementation for Dictionaries will encode each key/value as it's own property/field. The ArrowEncoder checks to see if the Type that is being encoded is a Dictionary and if so, sets the member variable modForIndex to 2. This ensures that the RecordBatch will contain only two columns (one for the key and value). ### Are these changes tested? Yes, unit tests has been added * GitHub Issue: #42021 Authored-by: Alva Bandy <[email protected]> Signed-off-by: Sutou Kouhei <[email protected]> --- Arrow/Sources/Arrow/ArrowEncoder.swift | 456 ++++++++++++++++++++++++++++++ Arrow/Tests/ArrowTests/CodableTests.swift | 108 ++++++- 2 files changed, 563 insertions(+), 1 deletion(-) diff --git a/Arrow/Sources/Arrow/ArrowEncoder.swift b/Arrow/Sources/Arrow/ArrowEncoder.swift new file mode 100644 index 0000000..8c72c04 --- /dev/null +++ b/Arrow/Sources/Arrow/ArrowEncoder.swift @@ -0,0 +1,456 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +import Foundation + +public class ArrowEncoder: Encoder { + public private(set) var builders = [String: ArrowArrayHolderBuilder]() + private var byIndex = [String]() + public var codingPath: [CodingKey] = [] + public var userInfo: [CodingUserInfoKey: Any] = [:] + var errorMsg: String? + // this is used for Dictionary types. A dictionary type + // will give each key and value there own index so instead + // of having a 2 column RecordBatch you would have + // 2 * length(dictionary) column RecordBatch. Which would not + // be the expected output. + var modForIndex: Int? + + public init() {} + + public init(_ builders: [String: ArrowArrayHolderBuilder], byIndex: [String]) { + self.builders = builders + self.byIndex = byIndex + } + + public static func encode<T: Encodable>(_ data: T) throws -> RecordBatch { + let encoder = try loadEncoder(data) + try data.encode(to: encoder) + return try encoder.finish() + } + + public static func encode<T: Encodable>(_ rows: [T]) throws -> RecordBatch? { + if rows.isEmpty { + return nil + } + + let encoder = try loadEncoder(rows[0]) + for row in rows { + try row.encode(to: encoder) + } + + return try encoder.finish() + } + + static func loadEncoder<T>(_ data: T) throws -> ArrowEncoder { + // this will check if T is a simple built in type + // (UInt, Int, Int8, String, Date, etc...). + if ArrowArrayBuilders.isValidBuilderType(T.self) { + let builders = ["col0": try ArrowArrayBuilders.loadBuilder(T.self)] + return ArrowEncoder(builders, byIndex: ["col0"]) + } else { + let encoder = ArrowEncoder() + if data is [AnyHashable: Any] { + encoder.modForIndex = 2 + } + + return encoder + } + } + + public func finish() throws -> RecordBatch { + try throwIfInvalid() + let batchBuilder = RecordBatch.Builder() + for key in byIndex { + batchBuilder.addColumn(key, arrowArray: try builders[key]!.toHolder()) + } + + switch batchBuilder.finish() { + case .success(let rb): + return rb + case .failure(let error): + throw error + } + } + + public func container<Key>(keyedBy type: Key.Type) -> KeyedEncodingContainer<Key> where Key: CodingKey { + var container = ArrowKeyedEncoding<Key>(self) + container.codingPath = codingPath + return KeyedEncodingContainer(container) + } + + public func unkeyedContainer() -> UnkeyedEncodingContainer { + return ArrowUnkeyedEncoding(self, codingPath: self.codingPath) + } + + public func singleValueContainer() -> SingleValueEncodingContainer { + return ArrowSingleValueEncoding(self, codingPath: codingPath) + } + + func doEncodeNil(key: CodingKey) throws { + try throwIfInvalid() + guard let builder = builders[key.stringValue] else { + throw ArrowError.invalid("Column not found for key: \(key)") + } + + builder.appendAny(nil) + } + + // This is required by the keyed and unkeyed encoders as columns are + // added when the first row of the data is encoded. This is done due + // to limitations in the Swifts Mirror API (ex: it is unable to correctly + // find the type for String? in [Int: String?]) + @discardableResult + func ensureColumnExists<T>(_ value: T, key: String) throws -> ArrowArrayHolderBuilder { + try throwIfInvalid() + var builder = builders[key] + if builder == nil { + builder = try ArrowArrayBuilders.loadBuilder(T.self) + builders[key] = builder + byIndex.append(key) + } + + return builder! + } + + func getIndex(_ index: Int) -> Int { + return self.modForIndex == nil ? index : index % self.modForIndex! + } + + func doEncodeNil(_ keyIndex: Int) throws { + try throwIfInvalid() + let index = self.getIndex(keyIndex) + if index >= builders.count { + throw ArrowError.outOfBounds(index: Int64(index)) + } + + builders[byIndex[index]]!.appendAny(nil) + } + + func doEncode<T>(_ value: T, key: CodingKey) throws { + try throwIfInvalid() + let builder = try ensureColumnExists(value, key: key.stringValue) + builder.appendAny(value) + } + + func doEncode<T>(_ value: T, keyIndex: Int) throws { + try throwIfInvalid() + let index = self.getIndex(keyIndex) + if index >= builders.count { + if index == builders.count { + try ensureColumnExists(value, key: "col\(index)") + } else { + throw ArrowError.outOfBounds(index: Int64(index)) + } + } + + builders[byIndex[index]]!.appendAny(value) + } + + func throwIfInvalid() throws { + if let errorMsg = self.errorMsg { + throw ArrowError.invalid(errorMsg) + } + } +} + +private struct ArrowKeyedEncoding<Key: CodingKey>: KeyedEncodingContainerProtocol { + var codingPath: [CodingKey] = [] + let encoder: ArrowEncoder + init(_ encoder: ArrowEncoder) { + self.encoder = encoder + } + + // If this method is called on row 0 and the encoder is + // lazily bulding holders then this will produce an error + // as this method does not know what the underlying type + // is for the column. This method is not called for + // nullable types (String?, Int32?, Date?) and the workaround + // for this issue would be to predefine the builders for the + // encoder. (I have only encoutered this issue when allowing + // nullable types at the encode func level which is currently + // not allowed) + mutating func encodeNil(forKey key: Key) throws { + try encoder.doEncodeNil(key: key) + } + + mutating func doEncodeIf<T>(_ value: T?, forKey key: Key) throws { + if value == nil { + try encoder.ensureColumnExists(value, key: key.stringValue) + try encoder.doEncodeNil(key: key) + } else { + try encoder.doEncode(value, key: key) + } + } + + mutating func encode(_ value: Bool, forKey key: Key) throws { + try encoder.doEncode(value, key: key) + } + + mutating func encodeIfPresent(_ value: Bool?, forKey key: Key) throws { + try doEncodeIf(value, forKey: key) + } + + mutating func encode(_ value: String, forKey key: Key) throws { + try encoder.doEncode(value, key: key) + } + + mutating func encodeIfPresent(_ value: String?, forKey key: Key) throws { + try doEncodeIf(value, forKey: key) + } + + mutating func encode(_ value: Double, forKey key: Key) throws { + try encoder.doEncode(value, key: key) + } + + mutating func encodeIfPresent(_ value: Double?, forKey key: Key) throws { + try doEncodeIf(value, forKey: key) + } + + mutating func encode(_ value: Float, forKey key: Key) throws { + try encoder.doEncode(value, key: key) + } + + mutating func encodeIfPresent(_ value: Float?, forKey key: Key) throws { + try doEncodeIf(value, forKey: key) + } + + mutating func encode(_ value: Int, forKey key: Key) throws { + throw ArrowError.invalid( + "Int type is not supported (please use Int8, Int16, Int32 or Int64)") + } + + mutating func encodeIfPresent(_ value: Int?, forKey key: Key) throws { + throw ArrowError.invalid( + "Int type is not supported (please use Int8, Int16, Int32 or Int64)") + } + + mutating func encode(_ value: Int8, forKey key: Key) throws { + try encoder.doEncode(value, key: key) + } + + mutating func encodeIfPresent(_ value: Int8?, forKey key: Key) throws { + try doEncodeIf(value, forKey: key) + } + + mutating func encode(_ value: Int16, forKey key: Key) throws { + try encoder.doEncode(value, key: key) + } + + mutating func encodeIfPresent(_ value: Int16?, forKey key: Key) throws { + try doEncodeIf(value, forKey: key) + } + + mutating func encode(_ value: Int32, forKey key: Key) throws { + try encoder.doEncode(value, key: key) + } + + mutating func encodeIfPresent(_ value: Int32?, forKey key: Key) throws { + try doEncodeIf(value, forKey: key) + } + + mutating func encode(_ value: Int64, forKey key: Key) throws { + try encoder.doEncode(value, key: key) + } + + mutating func encodeIfPresent(_ value: Int64?, forKey key: Key) throws { + try doEncodeIf(value, forKey: key) + } + + mutating func encode(_ value: UInt, forKey key: Key) throws { + throw ArrowError.invalid( + "UInt type is not supported (please use UInt8, UInt16, UInt32 or UInt64)") + } + + mutating func encodeIfPresent(_ value: UInt?, forKey key: Key) throws { + throw ArrowError.invalid( + "UInt type is not supported (please use UInt8, UInt16, UInt32 or UInt64)") + } + + mutating func encode(_ value: UInt8, forKey key: Key) throws { + try encoder.doEncode(value, key: key) + } + + mutating func encodeIfPresent(_ value: UInt8?, forKey key: Key) throws { + try doEncodeIf(value, forKey: key) + } + + mutating func encode(_ value: UInt16, forKey key: Key) throws { + try encoder.doEncode(value, key: key) + } + + mutating func encodeIfPresent(_ value: UInt16?, forKey key: Key) throws { + try doEncodeIf(value, forKey: key) + } + + mutating func encode(_ value: UInt32, forKey key: Key) throws { + try encoder.doEncode(value, key: key) + } + + mutating func encodeIfPresent(_ value: UInt32?, forKey key: Key) throws { + try doEncodeIf(value, forKey: key) + } + + mutating func encode(_ value: UInt64, forKey key: Key) throws { + try encoder.doEncode(value, key: key) + } + + mutating func encodeIfPresent(_ value: UInt64?, forKey key: Key) throws { + try doEncodeIf(value, forKey: key) + } + + mutating func encode<T: Encodable>(_ value: T, forKey key: Key) throws { + if ArrowArrayBuilders.isValidBuilderType(T.self) { + try encoder.doEncode(value, key: key) + } else { + throw ArrowError.invalid("Type \(T.self) is currently not supported") + } + } + + mutating func encodeIfPresent<T>(_ value: T?, forKey key: Self.Key) throws where T: Encodable { + if ArrowArrayBuilders.isValidBuilderType(T?.self) { + try doEncodeIf(value, forKey: key) + } else { + throw ArrowError.invalid("Type \(T.self) is currently not supported") + } + } + + // nested container is currently not allowed. This method doesn't throw + // so setting an error mesg that will be throw by the encoder at the next + // method call that throws + mutating func nestedContainer<NestedKey: CodingKey>( + keyedBy keyType: NestedKey.Type, + forKey key: Key) -> KeyedEncodingContainer<NestedKey> { + self.encoder.errorMsg = "Nested decoding is currently not supported." + var container = ArrowKeyedEncoding<NestedKey>(self.encoder) + container.codingPath = codingPath + return KeyedEncodingContainer(container) + } + + // nested container is currently not allowed. This method doesn't throw + // so setting an error mesg that will be throw by the encoder at the next + // method call that throws + mutating func nestedUnkeyedContainer(forKey key: Key) -> UnkeyedEncodingContainer { + self.encoder.errorMsg = "Nested decoding is currently not supported." + return ArrowUnkeyedEncoding(self.encoder, codingPath: self.codingPath) + } + + // super encoding is currently not allowed. This method doesn't throw + // so setting an error mesg that will be throw by the encoder at the next + // method call that throws + mutating func superEncoder() -> Encoder { + self.encoder.errorMsg = "super encoding is currently not supported." + return self.encoder + } + + // super encoding is currently not allowed. This method doesn't throw + // so setting an error mesg that will be throw by the encoder at the next + // method call that throws + mutating func superEncoder(forKey key: Key) -> Encoder { + self.encoder.errorMsg = "super encoding is currently not supported." + return self.encoder + } +} + +private struct ArrowUnkeyedEncoding: UnkeyedEncodingContainer { + public private(set) var encoder: ArrowEncoder + var codingPath: [CodingKey] = [] + var currentIndex: Int + var count: Int = 0 + + init(_ encoder: ArrowEncoder, codingPath: [CodingKey], currentIndex: Int = 0) { + self.encoder = encoder + self.currentIndex = currentIndex + } + + mutating func increment() { + self.currentIndex += 1 + } + + // If this method is called on row 0 and the encoder is + // lazily bulding holders then this will produce an error + // as this method does not know what the underlying type + // is for the column. This method is not called for + // nullable types (String?, Int32?, Date?) and the workaround + // for this issue would be to predefine the builders for the + // encoder. (I have only encoutered this issue when allowing + // nullable types at the encode func level which is currently + // not allowed) + mutating func encodeNil() throws { + try encoder.doEncodeNil(self.currentIndex) + } + + mutating func encode<T>(_ value: T) throws where T: Encodable { + let type = T.self + if ArrowArrayBuilders.isValidBuilderType(type) { + defer {increment()} + return try self.encoder.doEncode(value, keyIndex: self.currentIndex) + } else { + throw ArrowError.invalid("Type \(type) is currently not supported") + } + } + + // nested container is currently not allowed. This method doesn't throw + // so setting an error mesg that will be throw by the encoder at the next + // method call that throws + mutating func nestedContainer<NestedKey>(keyedBy keyType: NestedKey.Type + ) -> KeyedEncodingContainer<NestedKey> where NestedKey: CodingKey { + self.encoder.errorMsg = "Nested decoding is currently not supported." + var container = ArrowKeyedEncoding<NestedKey>(self.encoder) + container.codingPath = codingPath + return KeyedEncodingContainer(container) + } + + // nested container is currently not allowed. This method doesn't throw + // so setting an error mesg that will be throw by the encoder at the next + // method call that throws + mutating func nestedUnkeyedContainer() -> UnkeyedEncodingContainer { + self.encoder.errorMsg = "Nested decoding is currently not supported." + return ArrowUnkeyedEncoding(self.encoder, codingPath: self.codingPath) + } + + // super encoding is currently not allowed. This method doesn't throw + // so setting an error mesg that will be throw by the encoder at the next + // method call that throws + mutating func superEncoder() -> Encoder { + self.encoder.errorMsg = "super encoding is currently not supported." + return self.encoder + } +} + +private struct ArrowSingleValueEncoding: SingleValueEncodingContainer { + public private(set) var encoder: ArrowEncoder + var codingPath: [CodingKey] = [] + + public init(_ encoder: ArrowEncoder, codingPath: [CodingKey]) { + self.encoder = encoder + self.codingPath = codingPath + } + + mutating func encodeNil() throws { + return try self.encoder.doEncodeNil(0) + } + + mutating func encode<T: Encodable>(_ value: T) throws { + if ArrowArrayBuilders.isValidBuilderType(T.self) { + return try self.encoder.doEncode(value, keyIndex: 0) + } else { + throw ArrowError.invalid("Type \(T.self) is currently not supported") + } + } +} +// swiftlint:disable:this file_length diff --git a/Arrow/Tests/ArrowTests/CodableTests.swift b/Arrow/Tests/ArrowTests/CodableTests.swift index 400faa9..a0c4e11 100644 --- a/Arrow/Tests/ArrowTests/CodableTests.swift +++ b/Arrow/Tests/ArrowTests/CodableTests.swift @@ -18,7 +18,7 @@ import XCTest @testable import Arrow -final class CodableTests: XCTestCase { +final class CodableTests: XCTestCase { // swiftlint:disable:this type_body_length public class TestClass: Codable { public var propBool: Bool public var propInt8: Int8 @@ -225,4 +225,110 @@ final class CodableTests: XCTestCase { throw err } } + + func getArrayValue<T>(_ rb: RecordBatch, colIndex: Int, rowIndex: UInt) -> T? { + let anyArray = rb.columns[colIndex].array as! AnyArray // swiftlint:disable:this force_cast + return anyArray.asAny(UInt(rowIndex)) as? T + } + + func testArrowKeyedEncoder() throws { // swiftlint:disable:this function_body_length + var infos = [TestClass]() + for index in 0..<10 { + let tClass = TestClass() + let offset = index * 12 + tClass.propBool = index % 2 == 0 + tClass.propInt8 = Int8(offset + 1) + tClass.propInt16 = Int16(offset + 2) + tClass.propInt32 = Int32(offset + 3) + tClass.propInt64 = Int64(offset + 4) + tClass.propUInt8 = UInt8(offset + 5) + tClass.propUInt16 = UInt16(offset + 6) + tClass.propUInt32 = UInt32(offset + 7) + tClass.propUInt64 = UInt64(offset + 8) + tClass.propFloat = Float(offset + 9) + tClass.propDouble = index % 2 == 0 ? Double(offset + 10) : nil + tClass.propString = "\(offset + 11)" + tClass.propDate = Date.now + infos.append(tClass) + } + + let rb = try ArrowEncoder.encode(infos)! + XCTAssertEqual(Int(rb.length), infos.count) + XCTAssertEqual(rb.columns.count, 13) + XCTAssertEqual(rb.columns[0].type.id, ArrowTypeId.boolean) + XCTAssertEqual(rb.columns[1].type.id, ArrowTypeId.int8) + XCTAssertEqual(rb.columns[2].type.id, ArrowTypeId.int16) + XCTAssertEqual(rb.columns[3].type.id, ArrowTypeId.int32) + XCTAssertEqual(rb.columns[4].type.id, ArrowTypeId.int64) + XCTAssertEqual(rb.columns[5].type.id, ArrowTypeId.uint8) + XCTAssertEqual(rb.columns[6].type.id, ArrowTypeId.uint16) + XCTAssertEqual(rb.columns[7].type.id, ArrowTypeId.uint32) + XCTAssertEqual(rb.columns[8].type.id, ArrowTypeId.uint64) + XCTAssertEqual(rb.columns[9].type.id, ArrowTypeId.float) + XCTAssertEqual(rb.columns[10].type.id, ArrowTypeId.double) + XCTAssertEqual(rb.columns[11].type.id, ArrowTypeId.string) + XCTAssertEqual(rb.columns[12].type.id, ArrowTypeId.date64) + for index in 0..<10 { + let offset = index * 12 + XCTAssertEqual(getArrayValue(rb, colIndex: 0, rowIndex: UInt(index)), index % 2 == 0) + XCTAssertEqual(getArrayValue(rb, colIndex: 1, rowIndex: UInt(index)), Int8(offset + 1)) + XCTAssertEqual(getArrayValue(rb, colIndex: 2, rowIndex: UInt(index)), Int16(offset + 2)) + XCTAssertEqual(getArrayValue(rb, colIndex: 3, rowIndex: UInt(index)), Int32(offset + 3)) + XCTAssertEqual(getArrayValue(rb, colIndex: 4, rowIndex: UInt(index)), Int64(offset + 4)) + XCTAssertEqual(getArrayValue(rb, colIndex: 5, rowIndex: UInt(index)), UInt8(offset + 5)) + XCTAssertEqual(getArrayValue(rb, colIndex: 6, rowIndex: UInt(index)), UInt16(offset + 6)) + XCTAssertEqual(getArrayValue(rb, colIndex: 7, rowIndex: UInt(index)), UInt32(offset + 7)) + XCTAssertEqual(getArrayValue(rb, colIndex: 8, rowIndex: UInt(index)), UInt64(offset + 8)) + XCTAssertEqual(getArrayValue(rb, colIndex: 9, rowIndex: UInt(index)), Float(offset + 9)) + if index % 2 == 0 { + XCTAssertEqual(getArrayValue(rb, colIndex: 10, rowIndex: UInt(index)), Double(offset + 10)) + } else { + XCTAssertEqual(getArrayValue(rb, colIndex: 10, rowIndex: UInt(index)), Double?(nil)) + } + + XCTAssertEqual(getArrayValue(rb, colIndex: 11, rowIndex: UInt(index)), String(offset + 11)) + } + } + + func testArrowUnkeyedEncoder() throws { + var testMap = [Int8: String?]() + for index in 0..<10 { + testMap[Int8(index)] = "test\(index)" + } + + let rb = try ArrowEncoder.encode(testMap) + XCTAssertEqual(Int(rb.length), testMap.count) + XCTAssertEqual(rb.columns.count, 2) + XCTAssertEqual(rb.columns[0].type.id, ArrowTypeId.int8) + XCTAssertEqual(rb.columns[1].type.id, ArrowTypeId.string) + for index in 0..<10 { + let key: Int8 = getArrayValue(rb, colIndex: 0, rowIndex: UInt(index))! + let value: String = getArrayValue(rb, colIndex: 1, rowIndex: UInt(index))! + XCTAssertEqual("test\(key)", value) + } + } + + func testArrowSingleEncoder() throws { + var intArray = [Int32?]() + for index in 0..<100 { + if index == 10 { + intArray.append(nil) + } else { + intArray.append(Int32(index)) + } + } + + let rb = try ArrowEncoder.encode(intArray)! + XCTAssertEqual(Int(rb.length), intArray.count) + XCTAssertEqual(rb.columns.count, 1) + XCTAssertEqual(rb.columns[0].type.id, ArrowTypeId.int32) + for index in 0..<100 { + if index == 10 { + let anyArray = rb.columns[0].array as! AnyArray // swiftlint:disable:this force_cast + XCTAssertNil(anyArray.asAny(UInt(index))) + } else { + XCTAssertEqual(getArrayValue(rb, colIndex: 0, rowIndex: UInt(index)), Int32(index)) + } + } + } }
