diff --git a/swift/Arrow/Sources/Arrow/ArrowArray.swift b/swift/Arrow/Sources/Arrow/ArrowArray.swift index 32b6ba1704511..745c22fd62800 100644 --- a/swift/Arrow/Sources/Arrow/ArrowArray.swift +++ b/swift/Arrow/Sources/Arrow/ArrowArray.swift @@ -14,14 +14,13 @@ // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. - import Foundation public protocol ArrowArrayHolder { var type: ArrowType {get} var length: UInt {get} var nullCount: UInt {get} - var array: Any {get} + var array: AnyArray {get} var data: ArrowData {get} var getBufferData: () -> [Data] {get} var getBufferDataSizes: () -> [Int] {get} @@ -29,11 +28,11 @@ public protocol ArrowArrayHolder { } public class ArrowArrayHolderImpl: ArrowArrayHolder { - public let array: Any public let data: ArrowData public let type: ArrowType public let length: UInt public let nullCount: UInt + public let array: AnyArray public let getBufferData: () -> [Data] public let getBufferDataSizes: () -> [Int] public let getArrowColumn: (ArrowField, [ArrowArrayHolder]) throws -> ArrowColumn @@ -73,6 +72,50 @@ public class ArrowArrayHolderImpl: ArrowArrayHolder { return ArrowColumn(field, chunked: ChunkedArrayHolder(try ChunkedArray(arrays))) } } + + public static func loadArray( // swiftlint:disable:this cyclomatic_complexity + _ arrowType: ArrowType, with: ArrowData) throws -> ArrowArrayHolder { + switch arrowType.id { + case .int8: + return ArrowArrayHolderImpl(FixedArray(with)) + case .int16: + return ArrowArrayHolderImpl(FixedArray(with)) + case .int32: + return ArrowArrayHolderImpl(FixedArray(with)) + case .int64: + return ArrowArrayHolderImpl(FixedArray(with)) + case .uint8: + return ArrowArrayHolderImpl(FixedArray(with)) + case .uint16: + return ArrowArrayHolderImpl(FixedArray(with)) + case .uint32: + return ArrowArrayHolderImpl(FixedArray(with)) + case .uint64: + return ArrowArrayHolderImpl(FixedArray(with)) + case .double: + return ArrowArrayHolderImpl(FixedArray(with)) + case .float: + return ArrowArrayHolderImpl(FixedArray(with)) + case .date32: + return ArrowArrayHolderImpl(Date32Array(with)) + case .date64: + return ArrowArrayHolderImpl(Date64Array(with)) + case .time32: + return ArrowArrayHolderImpl(Time32Array(with)) + case .time64: + return ArrowArrayHolderImpl(Time64Array(with)) + case .string: + return ArrowArrayHolderImpl(StringArray(with)) + case .boolean: + return ArrowArrayHolderImpl(BoolArray(with)) + case .binary: + return ArrowArrayHolderImpl(BinaryArray(with)) + case .strct: + return ArrowArrayHolderImpl(StructArray(with)) + default: + throw ArrowError.invalid("Array not found for type: \(arrowType)") + } + } } public class ArrowArray: AsString, AnyArray { @@ -221,10 +264,7 @@ public class BinaryArray: ArrowArray { } public override func asString(_ index: UInt) -> String { - if self[index] == nil { - return "" - } - + if self[index] == nil {return ""} let data = self[index]! if options.printAsHex { return data.hexEncodedString() @@ -233,3 +273,58 @@ public class BinaryArray: ArrowArray { } } } + +public class StructArray: ArrowArray<[Any?]> { + public private(set) var arrowFields: [ArrowArrayHolder]? + public required init(_ arrowData: ArrowData) { + super.init(arrowData) + } + + public func initialize() throws -> StructArray { + var fields = [ArrowArrayHolder]() + for child in arrowData.children { + fields.append(try ArrowArrayHolderImpl.loadArray(child.type, with: child)) + } + + self.arrowFields = fields + return self + } + + public override subscript(_ index: UInt) -> [Any?]? { + if self.arrowData.isNull(index) { + return nil + } + + if let fields = arrowFields { + var result = [Any?]() + for field in fields { + result.append(field.array.asAny(index)) + } + + return result + } + + return nil + } + + public override func asString(_ index: UInt) -> String { + if self.arrowData.isNull(index) { + return "" + } + + var output = "{" + if let fields = arrowFields { + for fieldIndex in 0.. [Data] {self.holder.getBufferData} public var getBufferDataSizes: () -> [Int] {self.holder.getBufferDataSizes} diff --git a/swift/Arrow/Sources/Arrow/ChunkedArray.swift b/swift/Arrow/Sources/Arrow/ChunkedArray.swift index c5ccfe4aec0e6..fb5734f64b6ba 100644 --- a/swift/Arrow/Sources/Arrow/ChunkedArray.swift +++ b/swift/Arrow/Sources/Arrow/ChunkedArray.swift @@ -18,6 +18,7 @@ import Foundation public protocol AnyArray { + var arrowData: ArrowData {get} func asAny(_ index: UInt) -> Any? var length: UInt {get} } diff --git a/swift/Arrow/Tests/ArrowTests/CodableTests.swift b/swift/Arrow/Tests/ArrowTests/CodableTests.swift index a0c4e111e4360..b8f389a5e0089 100644 --- a/swift/Arrow/Tests/ArrowTests/CodableTests.swift +++ b/swift/Arrow/Tests/ArrowTests/CodableTests.swift @@ -227,7 +227,7 @@ final class CodableTests: XCTestCase { // swiftlint:disable:this type_body_lengt } func getArrayValue(_ rb: RecordBatch, colIndex: Int, rowIndex: UInt) -> T? { - let anyArray = rb.columns[colIndex].array as! AnyArray // swiftlint:disable:this force_cast + let anyArray = rb.columns[colIndex].array return anyArray.asAny(UInt(rowIndex)) as? T } @@ -324,7 +324,7 @@ final class CodableTests: XCTestCase { // swiftlint:disable:this type_body_lengt 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 + let anyArray = rb.columns[0].array XCTAssertNil(anyArray.asAny(UInt(index))) } else { XCTAssertEqual(getArrayValue(rb, colIndex: 0, rowIndex: UInt(index)), Int32(index))