fea(Arrow)t: Implement ListArray and ListArrayBuilder with associated functionality

Author: Marco

Arrow/Sources/Arrow/ArrowArray.swift

@@ -113,6 +113,8 @@ public class ArrowArrayHolderImpl: ArrowArrayHolder {
             return try ArrowArrayHolderImpl(BinaryArray(with))
         case .strct:
             return try ArrowArrayHolderImpl(StructArray(with))
+        case .list:
+            return try ArrowArrayHolderImpl(ListArray(with))
         default:
             throw ArrowError.invalid("Array not found for type: \(arrowType)")
         }
@@ -325,3 +327,69 @@ public class StructArray: ArrowArray<[Any?]> {
         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 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 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
+    }
+}

Arrow/Sources/Arrow/ArrowArrayBuilder.swift

@@ -125,7 +125,7 @@ public class StructArrayBuilder: ArrowArrayBuilder<StructBufferBuilder, StructAr
     public init(_ fields: [ArrowField], builders: [any ArrowArrayHolderBuilder]) throws {
         self.fields = fields
         self.builders = builders
-        try super.init(ArrowNestedType(ArrowType.ArrowStruct, fields: fields))
+        try super.init(ArrowTypeStruct(ArrowType.ArrowStruct, fields: fields))
         self.bufferBuilder.initializeTypeInfo(fields)
     }
 
@@ -137,7 +137,7 @@ public class StructArrayBuilder: ArrowArrayBuilder<StructBufferBuilder, StructAr
         }
 
         self.builders = builders
-        try super.init(ArrowNestedType(ArrowType.ArrowStruct, fields: fields))
+        try super.init(ArrowTypeStruct(ArrowType.ArrowStruct, fields: fields))
     }
 
     public override func append(_ values: [Any?]?) {
@@ -168,6 +168,31 @@ public class StructArrayBuilder: ArrowArrayBuilder<StructBufferBuilder, StructAr
     }
 }
 
+public class ListArrayBuilder: ArrowArrayBuilder<ListBufferBuilder, ListArray> {
+    let valueBuilder: any ArrowArrayHolderBuilder
+
+    public override init(_ elementType: ArrowType) throws {
+        self.valueBuilder = try ArrowArrayBuilders.loadBuilder(arrowType: elementType)
+        try super.init(ArrowTypeList(elementType))
+    }
+
+    public override func append(_ values: [Any?]?) {
+        self.bufferBuilder.append(values)
+        if let vals = values {
+            for val in vals {
+                self.valueBuilder.appendAny(val)
+            }
+        }
+    }
+
+    public override func finish() throws -> ListArray {
+        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)
+    }
+}
+
 public class ArrowArrayBuilders {
     public static func loadBuilder( // swiftlint:disable:this cyclomatic_complexity
         _ builderType: Any.Type) throws -> ArrowArrayHolderBuilder {
@@ -279,6 +304,16 @@ public class ArrowArrayBuilders {
                 throw ArrowError.invalid("Expected arrow type for \(arrowType.id) not found")
             }
             return try Time64ArrayBuilder(timeType.unit)
+        case .list:
+            guard let listType = arrowType as? ArrowTypeList else {
+                throw ArrowError.invalid("Expected ArrowTypeList for \(arrowType.id)")
+            }
+            return try ListArrayBuilder(listType.elementType)
+        case .strct:
+            guard let structType = arrowType as? ArrowTypeStruct else {
+                throw ArrowError.invalid("Expected ArrowStructType for \(arrowType.id)")
+            }
+            return try StructArrayBuilder(structType.fields)
         default:
             throw ArrowError.unknownType("Builder not found for arrow type: \(arrowType.id)")
         }
@@ -338,4 +373,12 @@ public class ArrowArrayBuilders {
     public static func loadTime64ArrayBuilder(_ unit: ArrowTime64Unit) throws -> Time64ArrayBuilder {
         return try Time64ArrayBuilder(unit)
     }
+
+    public static func loadStructArrayBuilder(_ fields: [ArrowField]) throws -> StructArrayBuilder {
+        return try StructArrayBuilder(fields)
+    }
+
+    public static func loadListArrayBuilder(_ elementType: ArrowType) throws -> ListArrayBuilder {
+        return try ListArrayBuilder(elementType)
+    }
 }

Arrow/Sources/Arrow/ArrowBufferBuilder.swift

@@ -338,14 +338,14 @@ public class Date64BufferBuilder: AbstractWrapperBufferBuilder<Date, Int64> {
 
 public final class StructBufferBuilder: BaseBufferBuilder, ArrowBufferBuilder {
     public typealias ItemType = [Any?]
-    var info: ArrowNestedType?
+    var info: ArrowTypeStruct?
     public init() throws {
         let nulls = ArrowBuffer.createBuffer(0, size: UInt(MemoryLayout<UInt8>.stride))
         super.init(nulls)
     }
 
     public func initializeTypeInfo(_ fields: [ArrowField]) {
-        info = ArrowNestedType(ArrowType.ArrowStruct, fields: fields)
+        info = ArrowTypeStruct(ArrowType.ArrowStruct, fields: fields)
     }
 
     public func append(_ newValue: [Any?]?) {
@@ -379,3 +379,62 @@ public final class StructBufferBuilder: BaseBufferBuilder, ArrowBufferBuilder {
         return [nulls]
     }
 }
+
+public class ListBufferBuilder: BaseBufferBuilder, ArrowBufferBuilder {
+    public typealias ItemType = [Any?]
+    var offsets: ArrowBuffer
+
+    public required init() throws {
+        self.offsets = ArrowBuffer.createBuffer(1, size: UInt(MemoryLayout<Int32>.stride))
+        let nulls = ArrowBuffer.createBuffer(0, size: UInt(MemoryLayout<UInt8>.stride))
+        super.init(nulls)
+        self.offsets.rawPointer.storeBytes(of: Int32(0), as: Int32.self)
+    }
+
+    public func append(_ newValue: [Any?]?) {
+        let index = UInt(self.length)
+        self.length += 1
+
+        if length >= self.offsets.length {
+            self.resize(length + 1)
+        }
+
+        let offsetIndex = Int(index) * MemoryLayout<Int32>.stride
+        let currentOffset = self.offsets.rawPointer.advanced(by: offsetIndex).load(as: Int32.self)
+
+        if let vals = newValue {
+            BitUtility.setBit(index + self.offset, buffer: self.nulls)
+            let newOffset = currentOffset + Int32(vals.count)
+            self.offsets.rawPointer.advanced(by: offsetIndex + MemoryLayout<Int32>.stride).storeBytes(of: newOffset, as: Int32.self)
+        } else {
+            self.nullCount += 1
+            BitUtility.clearBit(index + self.offset, buffer: self.nulls)
+            self.offsets.rawPointer.advanced(by: offsetIndex + MemoryLayout<Int32>.stride).storeBytes(of: currentOffset, as: Int32.self)
+        }
+    }
+
+    public override func isNull(_ index: UInt) -> Bool {
+        return !BitUtility.isSet(index + self.offset, buffer: self.nulls)
+    }
+
+    public func resize(_ length: UInt) {
+        if length > self.offsets.length {
+            let resizeLength = resizeLength(self.offsets)
+            var offsets = ArrowBuffer.createBuffer(resizeLength, size: UInt(MemoryLayout<Int32>.size))
+            var nulls = ArrowBuffer.createBuffer(resizeLength/8 + 1, size: UInt(MemoryLayout<UInt8>.size))
+            ArrowBuffer.copyCurrent(self.offsets, to: &offsets, len: self.offsets.capacity)
+            ArrowBuffer.copyCurrent(self.nulls, to: &nulls, len: self.nulls.capacity)
+            self.offsets = offsets
+            self.nulls = nulls
+        }
+    }
+
+    public func finish() -> [ArrowBuffer] {
+        let length = self.length
+        var nulls = ArrowBuffer.createBuffer(length/8 + 1, size: UInt(MemoryLayout<UInt8>.size))
+        var offsets = ArrowBuffer.createBuffer(length + 1, size: UInt(MemoryLayout<Int32>.size))
+        ArrowBuffer.copyCurrent(self.nulls, to: &nulls, len: nulls.capacity)
+        ArrowBuffer.copyCurrent(self.offsets, to: &offsets, len: offsets.capacity)
+        return [nulls, offsets]
+    }
+}

Arrow/Sources/Arrow/ArrowReader.swift

@@ -116,6 +116,35 @@ public class ArrowReader { // swiftlint:disable:this type_body_length
                                rbLength: UInt(loadInfo.batchData.recordBatch.length))
     }
 
+    private func loadListData(_ loadInfo: DataLoadInfo, field: org_apache_arrow_flatbuf_Field) -> Result<ArrowArrayHolder, ArrowError> {
+        guard let node = loadInfo.batchData.nextNode() else {
+            return .failure(.invalid("Node not found"))
+        }
+
+        guard let nullBuffer = loadInfo.batchData.nextBuffer() else {
+            return .failure(.invalid("Null buffer not found"))
+        }
+
+        guard let offsetBuffer = loadInfo.batchData.nextBuffer() else {
+            return .failure(.invalid("Offset buffer not found"))
+        }
+
+        let nullLength = UInt(ceil(Double(node.length) / 8))
+        let arrowNullBuffer = makeBuffer(nullBuffer, fileData: loadInfo.fileData, length: nullLength, messageOffset: loadInfo.messageOffset)
+        let arrowOffsetBuffer = makeBuffer(offsetBuffer, fileData: loadInfo.fileData, length: UInt(node.length + 1), messageOffset: loadInfo.messageOffset)
+
+        guard field.childrenCount == 1, let childField = field.children(at: 0) else {
+            return .failure(.invalid("List must have exactly one child"))
+        }
+
+        switch loadField(loadInfo, field: childField) {
+        case .success(let childHolder):
+            return makeArrayHolder(field, buffers: [arrowNullBuffer, arrowOffsetBuffer], nullCount: UInt(node.nullCount), children: [childHolder.array.arrowData], rbLength: UInt(loadInfo.batchData.recordBatch.length))
+        case .failure(let error):
+            return .failure(error)
+        }
+    }
+
     private func loadPrimitiveData(
         _ loadInfo: DataLoadInfo,
         field: org_apache_arrow_flatbuf_Field)
@@ -178,12 +207,17 @@ public class ArrowReader { // swiftlint:disable:this type_body_length
         _ loadInfo: DataLoadInfo,
         field: org_apache_arrow_flatbuf_Field)
     -> Result<ArrowArrayHolder, ArrowError> {
-        if isNestedType(field.typeType) {
+        switch field.typeType {
+        case .struct_:
             return loadStructData(loadInfo, field: field)
-        } else if isFixedPrimitive(field.typeType) {
-            return loadPrimitiveData(loadInfo, field: field)
-        } else {
-            return loadVariableData(loadInfo, field: field)
+        case .list:
+            return loadListData(loadInfo, field: field)
+        default:
+            if isFixedPrimitive(field.typeType) {
+                return loadPrimitiveData(loadInfo, field: field)
+            } else {
+                return loadVariableData(loadInfo, field: field)
+            }
         }
     }
 

Arrow/Sources/Arrow/ArrowReaderHelper.swift

@@ -136,6 +136,23 @@ func makeStructHolder(
     }
 }
 
+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)))
+    } catch let error as ArrowError {
+        return .failure(error)
+    } catch {
+        return .failure(.unknownError("\(error)"))
+    }
+}
+
 func makeArrayHolder(
     _ field: org_apache_arrow_flatbuf_Field,
     buffers: [ArrowBuffer],
@@ -188,6 +205,8 @@ func makeArrayHolder( // swiftlint:disable:this cyclomatic_complexity
         return makeTimeHolder(field, buffers: buffers, nullCount: nullCount)
     case .strct:
         return makeStructHolder(field, buffers: buffers, nullCount: nullCount, children: children!, rbLength: rbLength)
+    case .list:
+        return makeListHolder(field, buffers: buffers, nullCount: nullCount, children: children!, rbLength: rbLength)
     default:
         return .failure(.unknownType("Type \(typeId) currently not supported"))
     }
@@ -210,15 +229,6 @@ func isFixedPrimitive(_ type: org_apache_arrow_flatbuf_Type_) -> Bool {
     }
 }
 
-func isNestedType(_ type: org_apache_arrow_flatbuf_Type_) -> Bool {
-    switch type {
-    case .struct_:
-        return true
-    default:
-        return false
-    }
-}
-
 func findArrowType( // swiftlint:disable:this cyclomatic_complexity function_body_length
     _ field: org_apache_arrow_flatbuf_Field) -> ArrowType {
     let type = field.typeType
@@ -271,7 +281,13 @@ func findArrowType( // swiftlint:disable:this cyclomatic_complexity function_bod
                 ArrowField(childField.name ?? "", type: childType, isNullable: childField.nullable))
         }
 
-        return ArrowNestedType(ArrowType.ArrowStruct, fields: fields)
+        return ArrowTypeStruct(ArrowType.ArrowStruct, fields: fields)
+    case .list:
+        guard field.childrenCount == 1, let childField = field.children(at: 0) else {
+            return ArrowType(ArrowType.ArrowUnknown)
+        }
+        let childType = findArrowType(childField)
+        return ArrowTypeList(childType)
     default:
         return ArrowType(ArrowType.ArrowUnknown)
     }