diff --git a/Sources/SparkConnect/ArrowArray.swift b/Sources/SparkConnect/ArrowArray.swift index 61a4e49..b95a896 100644 --- a/Sources/SparkConnect/ArrowArray.swift +++ b/Sources/SparkConnect/ArrowArray.swift @@ -120,7 +120,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(NestedArray(with)) default: throw ArrowError.invalid("Array not found for type: \(arrowType)") } @@ -395,16 +397,40 @@ public class BinaryArray: ArrowArray { } /// @nodoc -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)) - } - self.arrowFields = fields + switch arrowData.type.id { + case .list: + 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 for list type ID") + } + + self.children = [ + try ArrowArrayHolderImpl.loadArray( + listType.elementField.type, + with: arrowData.children[0] + ) + ] + + case .strct: + var fields = [ArrowArrayHolder]() + for child in arrowData.children { + fields.append(try ArrowArrayHolderImpl.loadArray(child.type, with: child)) + } + self.children = fields + + default: + throw ArrowError.invalid( + "NestedArray only supports list and struct types, got: \(arrowData.type.id)") + } } public override subscript(_ index: UInt) -> [Any?]? { @@ -412,36 +438,105 @@ public class StructArray: ArrowArray<[Any?]> { return nil } - if let fields = arrowFields { + guard let children = self.children else { + return nil + } + + switch arrowData.type.id { + case .list: + guard let values = children.first else { return nil } + + let offsets = self.arrowData.buffers[1] + let offsetIndex = Int(index) * MemoryLayout.stride + + let startOffset = offsets.rawPointer.advanced(by: offsetIndex).load(as: Int32.self) + let endOffset = offsets.rawPointer.advanced(by: offsetIndex + MemoryLayout.stride) + .load(as: Int32.self) + + var items = [Any?]() + for i in startOffset.. String { - if self.arrowData.isNull(index) { - return "" - } + switch arrowData.type.id { + case .list: + if self.arrowData.isNull(index) { + return "null" + } + + guard let list = self[index] else { + return "null" + } - var output = "{" - if let fields = arrowFields { - for fieldIndex in 0.. 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 + + case .strct: + if self.arrowData.isNull(index) { + return "" + } + + var output = "{" + if let children = self.children { + for fieldIndex in 0.., } } -public class StructArrayBuilder: ArrowArrayBuilder { +public class StructArrayBuilder: ArrowArrayBuilder { let builders: [any ArrowArrayHolderBuilder] let fields: [ArrowField] 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) } @@ -153,7 +153,7 @@ public class StructArrayBuilder: ArrowArrayBuilder StructArray { + public override func finish() throws -> NestedArray { let buffers = self.bufferBuilder.finish() var childData = [ArrowData]() for builder in self.builders { @@ -180,11 +180,42 @@ public class StructArrayBuilder: ArrowArrayBuilder { + let valueBuilder: any ArrowArrayHolderBuilder + + public override init(_ arrowType: ArrowType) throws { + guard let listType = arrowType as? ArrowTypeList else { + throw ArrowError.invalid("Expected ArrowTypeList") + } + let arrowField = listType.elementField + self.valueBuilder = try ArrowArrayBuilders.loadBuilder(arrowType: arrowField.type) + try super.init(arrowType) + } + + 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 -> 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 NestedArray(arrowData) + } +} + public class ArrowArrayBuilders { public static func loadBuilder( // swiftlint:disable:this cyclomatic_complexity _ builderType: Any.Type @@ -304,6 +335,16 @@ public class ArrowArrayBuilders { throw ArrowError.invalid("Expected arrow type for \(arrowType.id) not found") } return try TimestampArrayBuilder(timestampType.unit) + case .strct: + guard let structType = arrowType as? ArrowTypeStruct else { + throw ArrowError.invalid("Expected ArrowStructType for \(arrowType.id)") + } + return try StructArrayBuilder(structType.fields) + case .list: + guard let listType = arrowType as? ArrowTypeList else { + throw ArrowError.invalid("Expected ArrowTypeList for \(arrowType.id)") + } + return try ListArrayBuilder(listType) default: throw ArrowError.unknownType("Builder not found for arrow type: \(arrowType.id)") } @@ -378,4 +419,12 @@ public class ArrowArrayBuilders { ) throws -> Decimal128ArrayBuilder { return try Decimal128ArrayBuilder(precision: precision, scale: scale) } + + public static func loadStructArrayBuilder(_ fields: [ArrowField]) throws -> StructArrayBuilder { + return try StructArrayBuilder(fields) + } + + public static func loadListArrayBuilder(_ listType: ArrowTypeList) throws -> ListArrayBuilder { + return try ListArrayBuilder(listType) + } } diff --git a/Sources/SparkConnect/ArrowBufferBuilder.swift b/Sources/SparkConnect/ArrowBufferBuilder.swift index b20e964..620c23a 100644 --- a/Sources/SparkConnect/ArrowBufferBuilder.swift +++ b/Sources/SparkConnect/ArrowBufferBuilder.swift @@ -343,20 +343,20 @@ public class Date64BufferBuilder: AbstractWrapperBufferBuilder { 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.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?]?) { let index = UInt(self.length) self.length += 1 - if length > self.nulls.length { + if self.length > self.nulls.length { self.resize(length) } @@ -385,3 +385,82 @@ 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.stride)) + let nulls = ArrowBuffer.createBuffer(0, size: UInt(MemoryLayout.stride)) + super.init(nulls) + self.offsets.rawPointer.storeBytes(of: Int32(0), as: Int32.self) + } + + public func append(_ count: Int) { + let index = UInt(self.length) + self.length += 1 + + if length >= self.offsets.length { + self.resize(length + 1) + } + + let offsetIndex = Int(index) * MemoryLayout.stride + let currentOffset = self.offsets.rawPointer.advanced(by: offsetIndex).load(as: Int32.self) + + BitUtility.setBit(index + self.offset, buffer: self.nulls) + let newOffset = currentOffset + Int32(count) + self.offsets.rawPointer.advanced(by: offsetIndex + MemoryLayout.stride).storeBytes( + of: newOffset, as: Int32.self) + } + + public func append(_ newValue: [Any?]?) { + let index = UInt(self.length) + self.length += 1 + + if self.length >= self.offsets.length { + self.resize(self.length + 1) + } + + let offsetIndex = Int(index) * MemoryLayout.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.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.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.size)) + var nulls = ArrowBuffer.createBuffer( + resizeLength / 8 + 1, size: UInt(MemoryLayout.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.size)) + var offsets = ArrowBuffer.createBuffer(length + 1, size: UInt(MemoryLayout.size)) + ArrowBuffer.copyCurrent(self.nulls, to: &nulls, len: nulls.capacity) + ArrowBuffer.copyCurrent(self.offsets, to: &offsets, len: offsets.capacity) + return [nulls, offsets] + } +} diff --git a/Sources/SparkConnect/ArrowReader.swift b/Sources/SparkConnect/ArrowReader.swift index f0699c2..c19a9e7 100644 --- a/Sources/SparkConnect/ArrowReader.swift +++ b/Sources/SparkConnect/ArrowReader.swift @@ -126,6 +126,44 @@ 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 + { + 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 @@ -204,6 +242,8 @@ public class ArrowReader { // swiftlint:disable:this type_body_length switch field.typeType { case .struct_: return loadStructData(loadInfo, field: field) + case .list: + return loadListData(loadInfo, field: field) default: if isFixedPrimitive(field.typeType) { return loadPrimitiveData(loadInfo, field: field) diff --git a/Sources/SparkConnect/ArrowReaderHelper.swift b/Sources/SparkConnect/ArrowReaderHelper.swift index bb0100a..ea397db 100644 --- a/Sources/SparkConnect/ArrowReaderHelper.swift +++ b/Sources/SparkConnect/ArrowReaderHelper.swift @@ -160,7 +160,7 @@ private func makeFixedHolder( } } -func makeStructHolder( +func makeNestedHolder( _ field: ArrowField, buffers: [ArrowBuffer], nullCount: UInt, @@ -170,9 +170,12 @@ func makeStructHolder( do { let arrowData = try ArrowData( field.type, - buffers: buffers, children: children, - nullCount: nullCount, length: rbLength) - return .success(ArrowArrayHolderImpl(try StructArray(arrowData))) + buffers: buffers, + children: children, + nullCount: nullCount, + length: rbLength + ) + return .success(ArrowArrayHolderImpl(try NestedArray(arrowData))) } catch let error as ArrowError { return .failure(error) } catch { @@ -236,7 +239,10 @@ func makeArrayHolder( // swiftlint:disable:this cyclomatic_complexity case .timestamp: return makeTimestampHolder(field, buffers: buffers, nullCount: nullCount) case .strct: - return makeStructHolder( + return makeNestedHolder( + field, buffers: buffers, nullCount: nullCount, children: children!, rbLength: rbLength) + case .list: + return makeNestedHolder( field, buffers: buffers, nullCount: nullCount, children: children!, rbLength: rbLength) default: return .failure(.unknownType("Type \(typeId) currently not supported")) @@ -345,7 +351,15 @@ func findArrowType( // swiftlint:disable:this cyclomatic_complexity function_bo 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) + let childFieldName = childField.name ?? "item" + return ArrowTypeList( + ArrowField(childFieldName, type: childType, isNullable: childField.nullable)) default: return ArrowType(ArrowType.ArrowUnknown) } diff --git a/Sources/SparkConnect/ArrowType.swift b/Sources/SparkConnect/ArrowType.swift index d7b773e..6ad39a7 100644 --- a/Sources/SparkConnect/ArrowType.swift +++ b/Sources/SparkConnect/ArrowType.swift @@ -190,7 +190,7 @@ public class ArrowTypeTimestamp: ArrowType { } /// @nodoc -public class ArrowNestedType: ArrowType { +public class ArrowTypeStruct: ArrowType { let fields: [ArrowField] public init(_ info: ArrowType.Info, fields: [ArrowField]) { self.fields = fields @@ -198,6 +198,19 @@ public class ArrowNestedType: ArrowType { } } +public class ArrowTypeList: ArrowType { + public let elementField: ArrowField + + public init(_ elementField: ArrowField) { + self.elementField = elementField + super.init(ArrowType.ArrowList) + } + + public convenience init(_ elementType: ArrowType, nullable: Bool = true) { + self.init(ArrowField("item", type: elementType, isNullable: nullable)) + } +} + /// @nodoc public class ArrowType { public private(set) var info: ArrowType.Info @@ -222,6 +235,7 @@ public class ArrowType { public static let ArrowTime64 = Info.timeInfo(ArrowTypeId.time64) public static let ArrowTimestamp = Info.timeInfo(ArrowTypeId.timestamp) public static let ArrowStruct = Info.complexInfo(ArrowTypeId.strct) + public static let ArrowList = Info.complexInfo(ArrowTypeId.list) public init(_ info: ArrowType.Info) { self.info = info @@ -355,7 +369,7 @@ public class ArrowType { return MemoryLayout.stride case .string: return MemoryLayout.stride - case .strct: + case .strct, .list: return 0 default: fatalError("Stride requested for unknown type: \(self)") @@ -412,6 +426,20 @@ public class ArrowType { return "z" case ArrowTypeId.string: return "u" + case ArrowTypeId.strct: + if let structType = self as? ArrowTypeStruct { + var format = "+s" + for field in structType.fields { + format += try field.type.cDataFormatId + } + return format + } + throw ArrowError.invalid("Invalid struct type") + case ArrowTypeId.list: + if let listType = self as? ArrowTypeList { + return "+l" + (try listType.elementField.type.cDataFormatId) + } + throw ArrowError.invalid("Invalid list type") default: throw ArrowError.notImplemented } diff --git a/Sources/SparkConnect/ArrowWriter.swift b/Sources/SparkConnect/ArrowWriter.swift index 9d44f9e..c7522a1 100644 --- a/Sources/SparkConnect/ArrowWriter.swift +++ b/Sources/SparkConnect/ArrowWriter.swift @@ -76,7 +76,7 @@ public class ArrowWriter { // swiftlint:disable:this type_body_length Offset, ArrowError > { var fieldsOffset: Offset? - if let nestedField = field.type as? ArrowNestedType { + if let nestedField = field.type as? ArrowTypeStruct { var offsets = [Offset]() for field in nestedField.fields { switch writeField(&fbb, field: field) { @@ -180,10 +180,11 @@ public class ArrowWriter { // swiftlint:disable:this type_body_length length: Int64(column.length), nullCount: Int64(column.nullCount)) offsets.append(fbb.create(struct: fieldNode)) - if let nestedType = column.type as? ArrowNestedType { - let structArray = column.array as? StructArray - writeFieldNodes( - nestedType.fields, columns: structArray!.arrowFields!, offsets: &offsets, fbb: &fbb) + if let nestedType = column.type as? ArrowTypeStruct { + let nestedArray = column.array as? NestedArray + if let nestedFields = nestedArray?.fields { + writeFieldNodes(nestedType.fields, columns: nestedFields, offsets: &offsets, fbb: &fbb) + } } } } @@ -204,11 +205,13 @@ public class ArrowWriter { // swiftlint:disable:this type_body_length offset: Int64(bufferOffset), length: Int64(bufferDataSize)) buffers.append(buffer) bufferOffset += bufferDataSize - if let nestedType = column.type as? ArrowNestedType { - let structArray = column.array as? StructArray - writeBufferInfo( - nestedType.fields, columns: structArray!.arrowFields!, - bufferOffset: &bufferOffset, buffers: &buffers, fbb: &fbb) + if let nestedType = column.type as? ArrowTypeStruct { + let nestedArray = column.array as? NestedArray + if let nestedFields = nestedArray?.fields { + writeBufferInfo( + nestedType.fields, columns: nestedFields, + bufferOffset: &bufferOffset, buffers: &buffers, fbb: &fbb) + } } } } @@ -267,13 +270,15 @@ public class ArrowWriter { // swiftlint:disable:this type_body_length for var bufferData in colBufferData { addPadForAlignment(&bufferData) writer.append(bufferData) - if let nestedType = column.type as? ArrowNestedType { - guard let structArray = column.array as? StructArray else { + if let nestedType = column.type as? ArrowTypeStruct { + 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!) + &writer, fields: nestedType.fields, columns: nestedFields) { case .success: continue diff --git a/Sources/SparkConnect/ProtoUtil.swift b/Sources/SparkConnect/ProtoUtil.swift index f890d07..6024315 100644 --- a/Sources/SparkConnect/ProtoUtil.swift +++ b/Sources/SparkConnect/ProtoUtil.swift @@ -96,7 +96,14 @@ func fromProto( // swiftlint:disable:this cyclomatic_complexity function_body_l children.append(fromProto(field: childField)) } - arrowType = ArrowNestedType(ArrowType.ArrowStruct, fields: children) + arrowType = ArrowTypeStruct(ArrowType.ArrowStruct, fields: children) + case .list: + guard field.childrenCount == 1, let childField = field.children(at: 0) else { + arrowType = ArrowType(ArrowType.ArrowUnknown) + break + } + let childArrowField = fromProto(field: childField) + arrowType = ArrowTypeList(childArrowField) default: arrowType = ArrowType(ArrowType.ArrowUnknown) }