Replacing ListArray and StructArray with NestedArray

Author: Marco

Arrow/Sources/Arrow/ArrowArray.swift

@@ -114,9 +114,9 @@ public class ArrowArrayHolderImpl: ArrowArrayHolder {
         case .binary:
             return try ArrowArrayHolderImpl(BinaryArray(with))
         case .strct:
-            return try ArrowArrayHolderImpl(StructArray(with))
+            return try ArrowArrayHolderImpl(NestedArray(with))
         case .list:
-            return try ArrowArrayHolderImpl(ListArray(with))
+            return try ArrowArrayHolderImpl(NestedArray(with))
         default:
             throw ArrowError.invalid("Array not found for type: \(arrowType)")
         }
@@ -357,119 +357,121 @@ public class BinaryArray: ArrowArray<Data> {
     }
 }
 
-public class StructArray: ArrowArray<[Any?]> {
-    public private(set) var arrowFields: [ArrowArrayHolder]?
+public class NestedArray: ArrowArray<[Any?]> {
+    private var children: [ArrowArrayHolder]?
+    
     public required init(_ arrowData: ArrowData) throws {
         try super.init(arrowData)
-        var fields = [ArrowArrayHolder]()
-        for child in arrowData.children {
-            fields.append(try ArrowArrayHolderImpl.loadArray(child.type, with: child))
+        
+        if let listType = arrowData.type as? ArrowTypeList {
+            guard arrowData.children.count == 1 else {
+                throw ArrowError.invalid("List array must have exactly one child")
+            }
+            
+            self.children = [try ArrowArrayHolderImpl.loadArray(
+                listType.elementType,
+                with: arrowData.children[0]
+            )]
+        } else {
+            var fields = [ArrowArrayHolder]()
+            for child in arrowData.children {
+                fields.append(try ArrowArrayHolderImpl.loadArray(child.type, with: child))
+            }
+            self.children = fields
         }
-
-        self.arrowFields = fields
     }
-
+    
     public override subscript(_ index: UInt) -> [Any?]? {
         if self.arrowData.isNull(index) {
             return nil
         }
-
-        if let fields = arrowFields {
+        
+        guard let children = self.children else {
+            return nil
+        }
+        
+        if arrowData.type is ArrowTypeList {
+            guard let values = children.first else { return nil }
+            
+            let offsets = self.arrowData.buffers[1]
+            let offsetIndex = Int(index) * MemoryLayout<Int32>.stride
+            
+            let startOffset = offsets.rawPointer.advanced(by: offsetIndex).load(as: Int32.self)
+            let endOffset = offsets.rawPointer.advanced(by: offsetIndex + MemoryLayout<Int32>.stride).load(as: Int32.self)
+            
+            var items = [Any?]()
+            for i in startOffset..<endOffset {
+                items.append(values.array.asAny(UInt(i)))
+            }
+            
+            return items
+        } else {
             var result = [Any?]()
-            for field in fields {
+            for field in children {
                 result.append(field.array.asAny(index))
             }
-
             return result
         }
-
-        return nil
     }
-
+    
     public override func asString(_ index: UInt) -> String {
+        let isListType = arrowData.type is ArrowTypeList
+        
         if self.arrowData.isNull(index) {
-            return ""
+            return isListType ? "null" : ""
         }
-
-        var output = "{"
-        if let fields = arrowFields {
-            for fieldIndex in 0..<fields.count {
-                let asStr = fields[fieldIndex].array as? AsString
-                if fieldIndex == 0 {
-                    output.append("\(asStr!.asString(index))")
+        
+        if isListType {
+            guard let list = self[index] else {
+                return "null"
+            }
+            
+            var output = "["
+            for (i, item) in list.enumerated() {
+                if i > 0 {
+                    output.append(",")
+                }
+                
+                if item == nil {
+                    output.append("null")
+                } else if let asStringItem = item as? AsString {
+                    output.append(asStringItem.asString(0))
                 } else {
-                    output.append(",\(asStr!.asString(index))")
+                    output.append("\(item!)")
                 }
             }
+            output.append("]")
+            return output
+        } else {
+            var output = "{"
+            if let children = self.children {
+                for fieldIndex in 0..<children.count {
+                    let asStr = children[fieldIndex].array as? AsString
+                    if fieldIndex == 0 {
+                        output.append("\(asStr!.asString(index))")
+                    } else {
+                        output.append(",\(asStr!.asString(index))")
+                    }
+                }
+            }
+            output += "}"
+            return output
         }
-
-        output += "}"
-        return output
     }
-}
-
-public class ListArray: ArrowArray<[Any?]> {
-    public private(set) var values: ArrowArrayHolder?
-
-    public required init(_ arrowData: ArrowData) throws {
-        try super.init(arrowData)
-        guard arrowData.children.count == 1 else {
-            throw ArrowError.invalid("List array must have exactly one child")
-        }
-
-        guard let listType = arrowData.type as? ArrowTypeList else {
-            throw ArrowError.invalid("Expected ArrowTypeList")
-        }
-
-        self.values = try ArrowArrayHolderImpl.loadArray(
-            listType.elementType,
-            with: arrowData.children[0]
-        )
+    
+    public var isListArray: Bool {
+        return arrowData.type is ArrowTypeList
     }
-
-    public override subscript(_ index: UInt) -> [Any?]? {
-        guard let values = self.values else { return nil }
-
-        if self.arrowData.isNull(index) {
-            return nil
-        }
-
-        let offsets = self.arrowData.buffers[1]
-        let offsetIndex = Int(index) * MemoryLayout<Int32>.stride
-
-        let startOffset = offsets.rawPointer.advanced(by: offsetIndex).load(as: Int32.self)
-        let endOffset = offsets.rawPointer.advanced(by: offsetIndex + MemoryLayout<Int32>.stride).load(as: Int32.self)
-
-        var items = [Any?]()
-        for i in startOffset..<endOffset {
-            items.append(values.array.asAny(UInt(i)))
-        }
-
-        return items
+    
+    public var isStructArray: Bool {
+        return arrowData.type is ArrowTypeStruct
     }
-
-    public override func asString(_ index: UInt) -> String {
-        guard let list = self[index] else {
-            return "null"
-        }
-
-        var output = "["
-
-        for (i, item) in list.enumerated() {
-            if i > 0 {
-                output.append(",")
-            }
-
-            if item == nil {
-                output.append("null")
-            } else if let asStringItem = item as? AsString {
-                output.append(asStringItem.asString(0))
-            } else {
-                output.append("\(item!)")
-            }
-        }
-
-        output.append("]")
-        return output
+    
+    public var fields: [ArrowArrayHolder]? {
+        return isStructArray ? children : nil
+    }
+    
+    public var values: ArrowArrayHolder? {
+        return isListArray ? children?.first : nil
     }
 }

Arrow/Sources/Arrow/ArrowArrayBuilder.swift

@@ -125,7 +125,7 @@ public class TimestampArrayBuilder: ArrowArrayBuilder<FixedBufferBuilder<Int64>,
     }
 }
 
-public class StructArrayBuilder: ArrowArrayBuilder<StructBufferBuilder, StructArray> {
+public class StructArrayBuilder: ArrowArrayBuilder<StructBufferBuilder, NestedArray> {
     let builders: [any ArrowArrayHolderBuilder]
     let fields: [ArrowField]
     public init(_ fields: [ArrowField], builders: [any ArrowArrayHolderBuilder]) throws {
@@ -159,7 +159,7 @@ public class StructArrayBuilder: ArrowArrayBuilder<StructBufferBuilder, StructAr
         }
     }
 
-    public override func finish() throws -> StructArray {
+    public override func finish() throws -> NestedArray {
         let buffers = self.bufferBuilder.finish()
         var childData = [ArrowData]()
         for builder in self.builders {
@@ -169,12 +169,12 @@ public class StructArrayBuilder: ArrowArrayBuilder<StructBufferBuilder, StructAr
         let arrowData = try ArrowData(self.type, buffers: buffers,
                                       children: childData, nullCount: self.nullCount,
                                       length: self.length)
-        let structArray = try StructArray(arrowData)
+        let structArray = try NestedArray(arrowData)
         return structArray
     }
 }
 
-public class ListArrayBuilder: ArrowArrayBuilder<ListBufferBuilder, ListArray> {
+public class ListArrayBuilder: ArrowArrayBuilder<ListBufferBuilder, NestedArray> {
     let valueBuilder: any ArrowArrayHolderBuilder
 
     public override init(_ elementType: ArrowType) throws {
@@ -191,11 +191,11 @@ public class ListArrayBuilder: ArrowArrayBuilder<ListBufferBuilder, ListArray> {
         }
     }
 
-    public override func finish() throws -> ListArray {
+    public override func finish() throws -> NestedArray {
         let buffers = self.bufferBuilder.finish()
         let childData = try valueBuilder.toHolder().array.arrowData
         let arrowData = try ArrowData(self.type, buffers: buffers, children: [childData], nullCount: self.nullCount, length: self.length)
-        return try ListArray(arrowData)
+        return try NestedArray(arrowData)
     }
 }
 

Arrow/Sources/Arrow/ArrowReaderHelper.swift

@@ -135,35 +135,22 @@ private func makeFixedHolder<T>(
     }
 }
 
-func makeStructHolder(
+func makeNestedHolder(
     _ field: ArrowField,
     buffers: [ArrowBuffer],
     nullCount: UInt,
     children: [ArrowData],
     rbLength: UInt
 ) -> Result<ArrowArrayHolder, ArrowError> {
     do {
-        let arrowData = try ArrowData(field.type,
-                                      buffers: buffers, children: children,
-                                      nullCount: nullCount, length: rbLength)
-        return .success(ArrowArrayHolderImpl(try StructArray(arrowData)))
-    } catch let error as ArrowError {
-        return .failure(error)
-    } catch {
-        return .failure(.unknownError("\(error)"))
-    }
-}
-
-func makeListHolder(
-    _ field: ArrowField,
-    buffers: [ArrowBuffer],
-    nullCount: UInt,
-    children: [ArrowData],
-    rbLength: UInt
-) -> Result<ArrowArrayHolder, ArrowError> {
-    do {
-        let arrowData = try ArrowData(field.type, buffers: buffers, children: children, nullCount: nullCount, length: rbLength)
-        return .success(ArrowArrayHolderImpl(try ListArray(arrowData)))
+        let arrowData = try ArrowData(
+            field.type,
+            buffers: buffers,
+            children: children,
+            nullCount: nullCount,
+            length: rbLength
+        )
+        return .success(ArrowArrayHolderImpl(try NestedArray(arrowData)))
     } catch let error as ArrowError {
         return .failure(error)
     } catch {
@@ -224,9 +211,9 @@ func makeArrayHolder( // swiftlint:disable:this cyclomatic_complexity
     case .timestamp:
         return makeTimestampHolder(field, buffers: buffers, nullCount: nullCount)
     case .strct:
-        return makeStructHolder(field, buffers: buffers, nullCount: nullCount, children: children!, rbLength: rbLength)
+        return makeNestedHolder(field, buffers: buffers, nullCount: nullCount, children: children!, rbLength: rbLength)
     case .list:
-        return makeListHolder(field, buffers: buffers, nullCount: nullCount, children: children!, rbLength: rbLength)
+        return makeNestedHolder(field, buffers: buffers, nullCount: nullCount, children: children!, rbLength: rbLength)
     default:
         return .failure(.unknownType("Type \(typeId) currently not supported"))
     }

Arrow/Sources/Arrow/ArrowWriter.swift

@@ -170,8 +170,10 @@ public class ArrowWriter { // swiftlint:disable:this type_body_length
                                                    nullCount: Int64(column.nullCount))
             offsets.append(fbb.create(struct: fieldNode))
             if let nestedType = column.type as? ArrowTypeStruct {
-                let structArray = column.array as? StructArray
-                writeFieldNodes(nestedType.fields, columns: structArray!.arrowFields!, offsets: &offsets, fbb: &fbb)
+                let nestedArray = column.array as? NestedArray
+                if let nestedFields = nestedArray?.fields {
+                    writeFieldNodes(nestedType.fields, columns: nestedFields, offsets: &offsets, fbb: &fbb)
+                }
             }
         }
     }
@@ -190,9 +192,11 @@ public class ArrowWriter { // swiftlint:disable:this type_body_length
                 buffers.append(buffer)
                 bufferOffset += bufferDataSize
                 if let nestedType = column.type as? ArrowTypeStruct {
-                    let structArray = column.array as? StructArray
-                    writeBufferInfo(nestedType.fields, columns: structArray!.arrowFields!,
-                                    bufferOffset: &bufferOffset, buffers: &buffers, fbb: &fbb)
+                    let nestedArray = column.array as? NestedArray
+                    if let nestedFields = nestedArray?.fields {
+                        writeBufferInfo(nestedType.fields, columns: nestedFields,
+                                        bufferOffset: &bufferOffset, buffers: &buffers, fbb: &fbb)
+                    }
                 }
             }
         }
@@ -236,22 +240,24 @@ public class ArrowWriter { // swiftlint:disable:this type_body_length
         return .success((fbb.data, Offset(offset: UInt32(fbb.data.count))))
     }
 
+    
     private func writeRecordBatchData(
         _ writer: inout DataWriter, fields: [ArrowField],
-        columns: [ArrowArrayHolder])
-    -> Result<Bool, ArrowError> {
+        columns: [ArrowArrayHolder]
+    ) -> Result<Bool, ArrowError> {
         for index in 0 ..< fields.count {
             let column = columns[index]
             let colBufferData = column.getBufferData()
             for var bufferData in colBufferData {
                 addPadForAlignment(&bufferData)
                 writer.append(bufferData)
                 if let nestedType = column.type as? ArrowTypeStruct {
-                    guard let structArray = column.array as? StructArray else {
+                    guard let nestedArray = column.array as? NestedArray,
+                          let nestedFields = nestedArray.fields else {
                         return .failure(.invalid("Struct type array expected for nested type"))
                     }
 
-                    switch writeRecordBatchData(&writer, fields: nestedType.fields, columns: structArray.arrowFields!) {
+                    switch writeRecordBatchData(&writer, fields: nestedType.fields, columns: nestedFields) {
                     case .success:
                         continue
                     case .failure(let error):

Arrow/Tests/ArrowTests/ArrayTests.swift

@@ -320,8 +320,8 @@ final class ArrayTests: XCTestCase { // swiftlint:disable:this type_body_length
         let structArray = try structBuilder.finish()
         XCTAssertEqual(structArray.length, 3)
         XCTAssertNil(structArray[1])
-        XCTAssertEqual(structArray.arrowFields![0].length, 3)
-        XCTAssertNil(structArray.arrowFields![0].array.asAny(1))
+        XCTAssertEqual(structArray.fields![0].length, 3)
+        XCTAssertNil(structArray.fields![0].array.asAny(1))
         XCTAssertEqual(structArray[0]![STIndex.bool.rawValue] as? Bool, true)
         XCTAssertEqual(structArray[0]![STIndex.int8.rawValue] as? Int8, 1)
         XCTAssertEqual(structArray[0]![STIndex.int16.rawValue] as? Int16, 2)