diff --git a/README.md b/README.md index 8a0a9c9..ef7d9b6 100644 --- a/README.md +++ b/README.md @@ -6,6 +6,104 @@ A Swift implementation of Apache Arrow, the universal columnar format for fast d This is a **work in progress**. Do not use in production. Progress is fast however, expect a beta in December. +## Array interface + +Arrow arrays are backed by a standard memory layout: +https://arrow.apache.org/docs/format/Columnar.html + +In Swift-Arrow, every array conforms to: + +```swift +public protocol ArrowArrayProtocol { + associatedtype ItemType + subscript(_ index: Int) -> ItemType? { get } + var offset: Int { get } + var length: Int { get } + func slice(offset: Int, length: Int) -> Self + func any(at index: Int) -> Any? +} +``` + +The in-memory contiguous buffers allow constant-time random access. + +Every Arrow array supports nullable elements. This is encoded as an optional bit-packed validity buffer aka null array aka bitfield. +In psuedocode, bitfield[index] == 0 means null or invalid, and bitfield[index] == 1 means not null or valid. +Fixed-width types are encoded back-to-back, with placeholder values for nulls. For example the array: + +```swift +let swiftArray: [Int8?] = [1, nil, 2, 3, nil, 4] +let arrayBuilder: ArrayBuilderFixedWidth = .init() +for value in swiftArray { + if let value { + arrayBuilder.append(value) + } else { + arrayBuilder.appendNull() + } +} +let arrowArray = arrayBuilder.finish() +for i in 0.. = .init() +for value in swiftArray { + if let value { + arrayBuilder.append(value) + } else { + arrayBuilder.appendNull() + } +} +let arrowArray = arrayBuilder.finish() +#expect(arrowArray[0] == "ab") +#expect(arrowArray[1] == nil) +#expect(arrowArray[2] == "c") +#expect(arrowArray[3] == "") +#expect(arrowArray[4] == ".") +``` + +would have an offsets array of array length + 1 integers: +`[0, 2, 2, 3, 3, 4]` + +This is a lookup into the value array, i.e.: + +```swift +let values: [UInt8] = [97, 98, 99, 46] +print(values[0..<2]) // [97, 98] +print(values[2..<2]) // [] +print(values[2..<3]) // [99] +print(values[3..<4]) // [46] +``` + +In practice, buffers can be any contingous storage. In Swift-Arrow, arrays created in memory are usually backed by pointers, whereas arrays loaded from IPC files are backed by memory-mapped `Data` instances. + +Arrays can be configured to use different buffer types, by specifying the types as +`public struct ArrowArrayVariable` + +this allows the buffer types to be user-specified, e.g.: +``` +typealias ArrowArrayUtf8 = ArrowArrayVariable< + FixedWidthBufferIPC, + VariableLengthBufferIPC +> +`` + + +## Relationship to Arrow-Swift + This project is based on Arrow-Swift, the official Swift implementation of Apache Arrow. The decision was made to at least temporarily operate independently of the Apache Software Foundation (ASF). Currently there are no active ASF maintaners with knowledge of Swift, and the only [Apache approved CI for Swift](https://github.com/apache/infrastructure-actions/blob/main/approved_patterns.yml) is [setup-swift which is unmaintained](https://github.com/swift-actions/setup-swift/issues), leading to intermittent CI failures. This has led to delays in much-needed fixes being implemented. The intention is to continue contributing to the official Apache-Swift repository, however changes can be iterated on more quickly here. diff --git a/Sources/Arrow/Array/Array.swift b/Sources/Arrow/Array/Array.swift index 8f867b2..0ac8c77 100644 --- a/Sources/Arrow/Array/Array.swift +++ b/Sources/Arrow/Array/Array.swift @@ -19,8 +19,10 @@ public protocol ArrowArrayProtocol { subscript(_ index: Int) -> ItemType? { get } var offset: Int { get } var length: Int { get } + var nullCount: Int { get } func slice(offset: Int, length: Int) -> Self func any(at index: Int) -> Any? + var bufferSizes: [Int] { get } } // This exists to support type-erased struct arrays. @@ -35,6 +37,8 @@ public struct ArrowArrayBoolean: ArrowArrayProtocol { public typealias ItemType = Bool public let offset: Int public let length: Int + public var bufferSizes: [Int] { [nullBuffer.length, valueBuffer.length] } + public var nullCount: Int { nullBuffer.nullCount } let nullBuffer: NullBuffer let valueBuffer: NullBuffer @@ -75,11 +79,12 @@ where ValueBuffer: FixedWidthBufferProtocol, ValueBuffer.ElementType: Numeric { - public typealias ItemType = ValueBuffer.ElementType - // public typealias ItemType = Element + public typealias ItemType = ValueBuffer.ElementType public let offset: Int public let length: Int + public var bufferSizes: [Int] { [nullBuffer.length, valueBuffer.length] } + public var nullCount: Int { nullBuffer.nullCount } let nullBuffer: NullBuffer let valueBuffer: ValueBuffer @@ -115,16 +120,20 @@ where } /// An Arrow array of variable-length types. -public struct ArrowArrayVariable: +public struct ArrowArrayVariable: ArrowArrayProtocol where - Element: VariableLength, OffsetsBuffer: FixedWidthBufferProtocol, - ValueBuffer: VariableLengthBufferProtocol + ValueBuffer: VariableLengthBufferProtocol, + ValueBuffer.ElementType: VariableLength { - public typealias ItemType = Element + public typealias ItemType = ValueBuffer.ElementType public let offset: Int public let length: Int + public var bufferSizes: [Int] { + [nullBuffer.length, offsetsBuffer.length, valueBuffer.length] + } + public var nullCount: Int { nullBuffer.nullCount } let nullBuffer: NullBuffer let offsetsBuffer: OffsetsBuffer let valueBuffer: ValueBuffer @@ -143,10 +152,8 @@ where self.valueBuffer = valueBuffer } - public subscript(index: Int) -> Element? { - + public subscript(index: Int) -> ValueBuffer.ElementType? { let offsetIndex = self.offset + index - if !self.nullBuffer.isSet(offsetIndex) { return nil } @@ -170,23 +177,18 @@ where } /// An Arrow array of `Date`s with a resolution of 1 day. -struct ArrowArrayDate32: ArrowArrayProtocol +public struct ArrowArrayDate32: ArrowArrayProtocol where ValueBuffer: FixedWidthBufferProtocol { - typealias ItemType = Date - + public typealias ItemType = Date + public var bufferSizes: [Int] { array.bufferSizes } + public var nullCount: Int { array.nullCount } + public var offset: Int { array.offset } + public var length: Int { array.length } let array: ArrowArrayFixed - var offset: Int { - array.offset - } - - var length: Int { - array.length - } - - subscript(index: Int) -> Date? { + public subscript(index: Int) -> Date? { precondition(index >= 0 && index < length, "Invalid index.") let offsetIndex = self.offset + index let days: Int32? = array[offsetIndex] @@ -197,30 +199,25 @@ where } } - func slice(offset: Int, length: Int) -> Self { + public func slice(offset: Int, length: Int) -> Self { let internalSlice = array.slice(offset: offset, length: length) return .init(array: internalSlice) } } /// An Arrow array of `Date`s with a resolution of 1 second. -struct ArrowArrayDate64: ArrowArrayProtocol +public struct ArrowArrayDate64: ArrowArrayProtocol where ValueBuffer: FixedWidthBufferProtocol { - typealias ItemType = Date - + public typealias ItemType = Date + public var bufferSizes: [Int] { array.bufferSizes } + public var nullCount: Int { array.nullCount } + public var offset: Int { array.offset } + public var length: Int { array.length } let array: ArrowArrayFixed - var offset: Int { - array.offset - } - - var length: Int { - array.length - } - - subscript(index: Int) -> Date? { + public subscript(index: Int) -> Date? { precondition(index >= 0 && index < length, "Invalid index.") let offsetIndex = self.offset + index let milliseconds: Int64? = array[offsetIndex] @@ -231,26 +228,44 @@ where } } - func slice(offset: Int, length: Int) -> Self { + public func slice(offset: Int, length: Int) -> Self { let internalSlice = array.slice(offset: offset, length: length) return .init(array: internalSlice) } } -/// An Arrow list array which may be nested arbitrarily. -struct ArrowListArray: ArrowArrayProtocol +/// A strongly-typed Arrow list array which may be nested arbitrarily. +public struct ArrowListArray: ArrowArrayProtocol where + OffsetsBuffer: FixedWidthBufferProtocol, Element: ArrowArrayProtocol { - typealias ItemType = Element - - let offset: Int - let length: Int + public typealias ItemType = Element + public let offset: Int + public let length: Int + public var bufferSizes: [Int] { + [nullBuffer.length, offsetsBuffer.length, values.length] + } + public var nullCount: Int { nullBuffer.nullCount } let nullBuffer: NullBuffer - let offsetsBuffer: FixedWidthBuffer + let offsetsBuffer: OffsetsBuffer let values: Element - subscript(index: Int) -> Element? { + public init( + offset: Int = 0, + length: Int, + nullBuffer: NullBuffer, + offsetsBuffer: OffsetsBuffer, + values: Element + ) { + self.offset = offset + self.length = length + self.nullBuffer = nullBuffer + self.offsetsBuffer = offsetsBuffer + self.values = values + } + + public subscript(index: Int) -> Element? { precondition(index >= 0 && index < length, "Invalid index.") let offsetIndex = self.offset + index if !self.nullBuffer.isSet(offsetIndex) { @@ -263,7 +278,7 @@ where return values.slice(offset: Int(startIndex), length: Int(length)) } - func slice(offset: Int, length: Int) -> Self { + public func slice(offset: Int, length: Int) -> Self { .init( offset: self.offset + offset, length: length, @@ -274,14 +289,54 @@ where } } +/// A type-erased wrapper for an Arrow list array. +public struct AnyArrowListArray: ArrowArrayProtocol { + + public typealias ItemType = any ArrowArrayProtocol + public var bufferSizes: [Int] { + _base.bufferSizes + } + + private let _base: any ArrowArrayProtocol + private let _subscriptImpl: (Int) -> (any ArrowArrayProtocol)? + private let _sliceImpl: (Int, Int) -> AnyArrowListArray + + public let offset: Int + public let length: Int + public var nullCount: Int { _base.nullCount } + + public init( + _ list: ArrowListArray + ) + where + OffsetsBuffer: FixedWidthBufferProtocol, + Element: ArrowArrayProtocol + { + self._base = list + self.offset = list.offset + self.length = list.length + self._subscriptImpl = { list[$0] } + self._sliceImpl = { AnyArrowListArray(list.slice(offset: $0, length: $1)) } + } + + public subscript(index: Int) -> (any ArrowArrayProtocol)? { + _subscriptImpl(index) + } + + public func slice(offset: Int, length: Int) -> AnyArrowListArray { + _sliceImpl(offset, length) + } +} + /// An Arrow struct array. public struct ArrowStructArray: ArrowArrayProtocol { public typealias ItemType = [String: Any] - - let nullBuffer: NullBuffer public let offset: Int public let length: Int public let fields: [(name: String, array: any ArrowArrayProtocol)] + public var bufferSizes: [Int] { [nullBuffer.length] } + public var nullCount: Int { nullBuffer.nullCount } + let nullBuffer: NullBuffer public init( offset: Int = 0, diff --git a/Sources/Arrow/Array/Builder.swift b/Sources/Arrow/Array/Builder.swift index 45617d9..5470b7d 100644 --- a/Sources/Arrow/Array/Builder.swift +++ b/Sources/Arrow/Array/Builder.swift @@ -35,15 +35,15 @@ extension AnyArrayBuilder { } /// A builder for Arrow arrays using the three-valued logical model (true / false / null). -class ArrayBuilderBoolean: AnyArrayBuilder { +public class ArrayBuilderBoolean: AnyArrayBuilder { - typealias ArrayType = ArrowArrayBoolean + public typealias ArrayType = ArrowArrayBoolean var length: Int let nullBuilder: NullBufferBuilder let valueBuilder: NullBufferBuilder - init() { + public init() { self.length = 0 self.nullBuilder = NullBufferBuilder() self.valueBuilder = NullBufferBuilder() @@ -55,7 +55,7 @@ class ArrayBuilderBoolean: AnyArrayBuilder { valueBuilder.appendValid(value) } - func appendNull() { + public func appendNull() { length += 1 nullBuilder.appendValid(false) } @@ -74,15 +74,15 @@ class ArrayBuilderBoolean: AnyArrayBuilder { } /// A builder for Arrow arrays holding fixed-width types. -class ArrayBuilderFixedWidth: AnyArrayBuilder { +public class ArrayBuilderFixedWidth: AnyArrayBuilder { - typealias ArrayType = ArrowArrayFixed> + public typealias ArrayType = ArrowArrayFixed> - var length: Int + public var length: Int let nullBuilder: NullBufferBuilder let valueBuilder: FixedWidthBufferBuilder - init() { + public init() { self.length = 0 self.nullBuilder = NullBufferBuilder() self.valueBuilder = FixedWidthBufferBuilder() @@ -94,7 +94,7 @@ class ArrayBuilderFixedWidth: AnyArrayBuilder { valueBuilder.append(value) } - func appendNull() { + public func appendNull() { length += 1 nullBuilder.appendValid(false) valueBuilder.append(T.zero) @@ -114,9 +114,11 @@ class ArrayBuilderFixedWidth: AnyArrayBuilder { } /// A builder for Arrow arrays holding variable length types. -class ArrayBuilderVariable: AnyArrayBuilder { - typealias ArrayType = ArrowArrayVariable< - Element, FixedWidthBuffer, VariableLengthTypeBuffer +public class ArrayBuilderVariableLength: + AnyArrayBuilder +{ + public typealias ArrayType = ArrowArrayVariable< + FixedWidthBuffer, VariableLengthTypeBuffer > var length: Int @@ -124,7 +126,7 @@ class ArrayBuilderVariable: AnyArrayBuilder { let offsetsBuilder: FixedWidthBufferBuilder let valueBuilder: VariableLengthTypeBufferBuilder - init() { + public init() { self.length = 0 self.nullBuilder = NullBufferBuilder() self.offsetsBuilder = FixedWidthBufferBuilder() @@ -149,7 +151,7 @@ class ArrayBuilderVariable: AnyArrayBuilder { offsetsBuilder.append(newOffset) } - func appendNull() { + public func appendNull() { length += 1 nullBuilder.appendValid(false) let newOffset = Int32(valueBuilder.length) @@ -171,18 +173,19 @@ class ArrayBuilderVariable: AnyArrayBuilder { } /// A builder for Arrow arrays holding `String` values. -typealias ArrayBuilderString = ArrayBuilderVariable +typealias ArrayBuilderString = ArrayBuilderVariableLength /// A builder for Arrow arrays holding `Data` values. -typealias ArrayBuilderBinary = ArrayBuilderVariable +typealias ArrayBuilderBinary = ArrayBuilderVariableLength /// A builder for Arrow arrays holding `Date`s with a resolution of one day. -struct ArrayBuilderDate32: AnyArrayBuilder { - typealias ArrayType = ArrowArrayDate32> - +public struct ArrayBuilderDate32: AnyArrayBuilder { + public typealias ArrayType = ArrowArrayDate32> let builder: ArrayBuilderFixedWidth = .init() - var length: Int { + public init() {} + + public var length: Int { builder.length } @@ -191,26 +194,27 @@ struct ArrayBuilderDate32: AnyArrayBuilder { self.builder.append(daysSinceEpoch) } - func appendNull() { + public func appendNull() { builder.appendNull() } - func finish() -> ArrayType { + public func finish() -> ArrayType { .init(array: builder.finish()) } } /// A builder for Arrow arrays holding `Date`s with a resolution of one day. -struct ArrayBuilderDate64: AnyArrayBuilder { - typealias ArrayType = ArrowArrayDate64> - +public struct ArrayBuilderDate64: AnyArrayBuilder { + public typealias ArrayType = ArrowArrayDate64> let builder: ArrayBuilderFixedWidth = .init() + public init() {} + var length: Int { builder.length } - func appendNull() { + public func appendNull() { self.builder.appendNull() } @@ -219,25 +223,25 @@ struct ArrayBuilderDate64: AnyArrayBuilder { self.builder.append(millisecondsSinceEpoch) } - func finish() -> ArrayType { + public func finish() -> ArrayType { .init(array: builder.finish()) } } /// A builder for Arrow arrays holding Time32 values. -typealias ArrayBuilderTime32 = ArrayBuilderFixedWidth +public typealias ArrayBuilderTime32 = ArrayBuilderFixedWidth /// A builder for Arrow arrays holding Time64 values. -typealias ArrayBuilderTime64 = ArrayBuilderFixedWidth +public typealias ArrayBuilderTime64 = ArrayBuilderFixedWidth /// A builder for Arrow arrays holding Timestamp values. -typealias ArrayBuilderTimestamp = ArrayBuilderFixedWidth +public typealias ArrayBuilderTimestamp = ArrayBuilderFixedWidth class ArrayBuilderList: AnyArrayBuilder { func append(_ value: T.ArrayType) {} - typealias ArrayType = ArrowListArray + typealias ArrayType = ArrowListArray> var length: Int let nullBuilder: NullBufferBuilder @@ -271,7 +275,7 @@ class ArrayBuilderList: AnyArrayBuilder { offsetsBuilder.append(Int32(valueBuilder.length)) } - func finish() -> ArrowListArray { + func finish() -> ArrayType { let nullBuffer = nullBuilder.finish() let offsetsBuffer = offsetsBuilder.finish() let valuesArray = valueBuilder.finish() @@ -319,7 +323,7 @@ class ArrayBuilderStruct: AnyArrayBuilder { func appendNull() { length += 1 nullBuilder.appendValid(false) - // Still need to append nulls to all child builders to keep lengths aligned + // Need to append nulls to all child builders to keep lengths aligned. for (_, builder) in fields { builder.appendNull() } diff --git a/Sources/Arrow/ArrowData.swift b/Sources/Arrow/ArrowData.swift index bf50734..bda48e4 100644 --- a/Sources/Arrow/ArrowData.swift +++ b/Sources/Arrow/ArrowData.swift @@ -53,7 +53,8 @@ public struct ArrowData { let buffer = buffers[0] let pointer = buffer.rawPointer.assumingMemoryBound(to: UInt8.self) return BitPackedNullBuffer( - length: Int(buffer.length), capacity: 0, ownsMemory: false, + length: Int(buffer.length), capacity: 0, valueCount: -1, nullCount: -1, + ownsMemory: false, buffer: pointer) } diff --git a/Sources/Arrow/ArrowDecoder.swift b/Sources/Arrow/ArrowDecoder.swift index c28be7b..e902e27 100644 --- a/Sources/Arrow/ArrowDecoder.swift +++ b/Sources/Arrow/ArrowDecoder.swift @@ -19,7 +19,7 @@ public class ArrowDecoder: Decoder { var singleRBCol: Int = 0 public var codingPath: [CodingKey] = [] public var userInfo: [CodingUserInfoKey: Any] = [:] - public let rb: RecordBatch + public let rb: RecordBatchX public let nameToCol: [String: AnyArrowArray] public let columns: [AnyArrowArray] @@ -32,7 +32,7 @@ public class ArrowDecoder: Decoder { self.rbIndex = decoder.rbIndex } - public init(_ rb: RecordBatch) { + public init(_ rb: RecordBatchX) { self.rb = rb var colMapping: [String: AnyArrowArray] = [:] var columns: [AnyArrowArray] = [] diff --git a/Sources/Arrow/ArrowEncoder.swift b/Sources/Arrow/ArrowEncoder.swift index f15573d..bf2dce6 100644 --- a/Sources/Arrow/ArrowEncoder.swift +++ b/Sources/Arrow/ArrowEncoder.swift @@ -34,13 +34,13 @@ public class ArrowEncoder: Encoder { self.byIndex = byIndex } - public static func encode(_ data: T) throws -> RecordBatch { + public static func encode(_ data: T) throws -> RecordBatchX { let encoder = try loadEncoder(data) try data.encode(to: encoder) return try encoder.finish() } - public static func encode(_ rows: [T]) throws -> RecordBatch? { + public static func encode(_ rows: [T]) throws -> RecordBatchX? { if rows.isEmpty { return nil } @@ -68,9 +68,9 @@ public class ArrowEncoder: Encoder { } } - public func finish() throws(ArrowError) -> RecordBatch { + public func finish() throws(ArrowError) -> RecordBatchX { try throwIfInvalid() - let batchBuilder = RecordBatch.Builder() + let batchBuilder = RecordBatchX.Builder() for key in byIndex { guard let builder = builders[key] else { throw .invalid("Missing builder for \(key)") diff --git a/Sources/Arrow/ArrowReader.swift b/Sources/Arrow/ArrowReaderX.swift similarity index 99% rename from Sources/Arrow/ArrowReader.swift rename to Sources/Arrow/ArrowReaderX.swift index 2fed70c..22b3668 100644 --- a/Sources/Arrow/ArrowReader.swift +++ b/Sources/Arrow/ArrowReaderX.swift @@ -19,7 +19,7 @@ import Foundation let fileMarker = Data("ARROW1".utf8) let continuationMarker = UInt32(0xFFFF_FFFF) -public struct ArrowReader: Sendable { +public struct ArrowReaderX: Sendable { func makeBuffer( _ buffer: FBuffer, @@ -84,7 +84,7 @@ public struct ArrowReader: Sendable { public class ArrowReaderResult { fileprivate var messageSchema: FSchema? public var schema: ArrowSchema? - public var batches: [RecordBatch] = [] + public var batches: [RecordBatchX] = [] } public init() {} @@ -301,7 +301,7 @@ public struct ArrowReader: Sendable { arrowSchema: ArrowSchema, data: Data, messageEndOffset: Int64 - ) -> Result { + ) -> Result { var columns: [AnyArrowArray] = [] let batchData = RecordBatchData(recordBatch, schema: schema) let loadInfo = DataLoadInfo( @@ -321,7 +321,7 @@ public struct ArrowReader: Sendable { return .failure(error) } } - return .success(RecordBatch(arrowSchema, columns: columns)) + return .success(RecordBatchX(arrowSchema, columns: columns)) } /// This is for reading the Arrow streaming format. diff --git a/Sources/Arrow/ArrowTable.swift b/Sources/Arrow/ArrowTable.swift index 82bcf9d..5ee70ce 100644 --- a/Sources/Arrow/ArrowTable.swift +++ b/Sources/Arrow/ArrowTable.swift @@ -53,7 +53,7 @@ public class ArrowTable { /// - Returns: An `ArrowResult` holding an `ArrowTable` on success, or an`ArrowError` /// on failure. public static func from( - recordBatches: [RecordBatch] + recordBatches: [RecordBatchX] ) -> Result { if recordBatches.isEmpty { return .failure(.arrayHasNoElements) @@ -206,7 +206,7 @@ public class ArrowTable { } } -public class RecordBatch { +public class RecordBatchX { public let schema: ArrowSchema public var columnCount: UInt { UInt(self.columns.count) } public let columns: [AnyArrowArray] @@ -261,7 +261,7 @@ public class RecordBatch { return self } - public func finish() -> Result { + public func finish() -> Result { if columns.count > 0 { let columnLength = columns[0].length for column in columns { @@ -282,7 +282,7 @@ public class RecordBatch { } } return .success( - RecordBatch(self.schemaBuilder.finish(), columns: self.columns) + RecordBatchX(self.schemaBuilder.finish(), columns: self.columns) ) } } diff --git a/Sources/Arrow/ArrowWriter.swift b/Sources/Arrow/ArrowWriterX.swift similarity index 98% rename from Sources/Arrow/ArrowWriter.swift rename to Sources/Arrow/ArrowWriterX.swift index d5f9bf7..35d818b 100644 --- a/Sources/Arrow/ArrowWriter.swift +++ b/Sources/Arrow/ArrowWriterX.swift @@ -56,12 +56,12 @@ public class ArrowWriter { public class Info { let type: FMessageHeader public let schema: ArrowSchema - public let batches: [RecordBatch] + public let batches: [RecordBatchX] init( _ type: FMessageHeader, schema: ArrowSchema, - batches: [RecordBatch] + batches: [RecordBatchX] ) { self.type = type self.schema = schema @@ -69,7 +69,7 @@ public class ArrowWriter { } convenience init(_ type: FMessageHeader, schema: ArrowSchema) { - self.init(type, schema: schema, batches: [RecordBatch]()) + self.init(type, schema: schema, batches: [RecordBatchX]()) } } @@ -143,7 +143,7 @@ public class ArrowWriter { private func writeRecordBatches( _ writer: inout DataWriter, - batches: [RecordBatch] + batches: [RecordBatchX] ) -> Result<[FBlock], ArrowError> { var rbBlocks: [FBlock] = .init() for batch in batches { @@ -249,7 +249,7 @@ public class ArrowWriter { } private func writeRecordBatch( - batch: RecordBatch + batch: RecordBatchX ) -> Result<(Data, Offset), ArrowError> { let schema = batch.schema var fbb = FlatBufferBuilder() @@ -482,7 +482,7 @@ public class ArrowWriter { return .success(true) } - public func toMessage(_ batch: RecordBatch) -> Result<[Data], ArrowError> { + public func toMessage(_ batch: RecordBatchX) -> Result<[Data], ArrowError> { var writer: any DataWriter = InMemDataWriter() switch writeRecordBatch(batch: batch) { case .success(let message): diff --git a/Sources/Arrow/Buffer/FixedBufferBuilder.swift b/Sources/Arrow/Buffer/FixedBufferBuilder.swift index 3fb2d9a..c346889 100644 --- a/Sources/Arrow/Buffer/FixedBufferBuilder.swift +++ b/Sources/Arrow/Buffer/FixedBufferBuilder.swift @@ -13,8 +13,8 @@ // limitations under the License. final class FixedWidthBufferBuilder { - var length: Int - var capacity: Int + var valueCount: Int + var valueCapacity: Int private var buffer: UnsafeMutablePointer private var ownsMemory: Bool private var bitOffset: Int8 = 0 @@ -22,32 +22,32 @@ final class FixedWidthBufferBuilder { init( minCapacity: Int = 1024 ) { - self.length = 0 + self.valueCount = 0 // Ensure at least 1 element capacity - self.capacity = max(1, minCapacity / MemoryLayout.size) - self.buffer = .allocate(capacity: capacity) + self.valueCapacity = max(1, minCapacity / MemoryLayout.size) + self.buffer = .allocate(capacity: valueCapacity) self.ownsMemory = true } func append(_ val: T) { - if length >= capacity { - var newCapacity = capacity * 2 - while length >= newCapacity { + if valueCount >= valueCapacity { + var newCapacity = valueCapacity * 2 + while valueCount >= newCapacity { newCapacity *= 2 } resize(to: newCapacity) } - buffer[length] = val - length += 1 + buffer[valueCount] = val + valueCount += 1 } private func resize(to newCapacity: Int) { - precondition(newCapacity > capacity) + precondition(newCapacity > valueCapacity) let newBuffer = UnsafeMutablePointer.allocate(capacity: newCapacity) - newBuffer.initialize(from: buffer, count: length) + newBuffer.initialize(from: buffer, count: valueCount) buffer.deallocate() buffer = newBuffer - capacity = newCapacity + valueCapacity = newCapacity } deinit { @@ -64,17 +64,18 @@ final class FixedWidthBufferBuilder { func finish() -> FixedWidthBuffer { precondition(ownsMemory, "Buffer already finished.") ownsMemory = false - let byteCount = length * MemoryLayout.size + let byteCount = valueCount * MemoryLayout.size let newCapacity = (byteCount + 63) & ~63 let newBuffer = UnsafeMutableRawPointer.allocate( byteCount: newCapacity, alignment: 64 ).bindMemory(to: T.self, capacity: newCapacity) - newBuffer.initialize(from: buffer, count: length) + newBuffer.initialize(from: buffer, count: valueCount) buffer.deallocate() return FixedWidthBuffer( - length: length, + length: byteCount, capacity: newCapacity, + valueCount: valueCount, ownsMemory: true, buffer: newBuffer ) diff --git a/Sources/Arrow/Buffer/FixedWidthBuffer.swift b/Sources/Arrow/Buffer/FixedWidthBuffer.swift index beffd68..b0f6c19 100644 --- a/Sources/Arrow/Buffer/FixedWidthBuffer.swift +++ b/Sources/Arrow/Buffer/FixedWidthBuffer.swift @@ -18,29 +18,36 @@ public protocol FixedWidthBufferProtocol { subscript(index: Int) -> ElementType { get } } -/// A buffer used in Arrow arrays that hold fixed-width types. -final class FixedWidthBuffer: FixedWidthBufferProtocol where T: Numeric { - - typealias ElementType = T +public protocol Int32BufferProtocol { + var length: Int { get } + subscript(index: Int) -> Int32 { get } +} - var length: Int +/// A buffer used in Arrow arrays that hold fixed-width types. +public final class FixedWidthBuffer: FixedWidthBufferProtocol +where T: Numeric { + public typealias ElementType = T + public var length: Int var capacity: Int + let valueCount: Int let ownsMemory: Bool var buffer: UnsafePointer init( length: Int, capacity: Int, + valueCount: Int, ownsMemory: Bool, buffer: UnsafePointer ) { self.length = length self.capacity = capacity + self.valueCount = valueCount self.ownsMemory = ownsMemory self.buffer = buffer } - subscript(index: Int) -> T { + public subscript(index: Int) -> T { buffer[index] } diff --git a/Sources/Arrow/Buffer/NullBuffer.swift b/Sources/Arrow/Buffer/NullBuffer.swift index 73b9fe1..8e4c31d 100644 --- a/Sources/Arrow/Buffer/NullBuffer.swift +++ b/Sources/Arrow/Buffer/NullBuffer.swift @@ -15,33 +15,39 @@ /// A type used to represent nulls and booleans in Arrow arrays. public protocol NullBuffer { var length: Int { get } + var valueCount: Int { get } + var nullCount: Int { get } func isSet(_ bit: Int) -> Bool } /// Represents an array with no nulls (all values valid). public struct AllValidNullBuffer: NullBuffer { - public let length: Int + public let valueCount: Int + public var length: Int { 0 } + public var nullCount: Int { 0 } - public init(length: Int) { - self.length = length + public init(valueCount: Int) { + self.valueCount = valueCount } public func isSet(_ bit: Int) -> Bool { - precondition(bit < length) + precondition(bit < valueCount) return true } } /// Represents an array with all nulls. public struct AllNullBuffer: NullBuffer { - public let length: Int + public let valueCount: Int + public var length: Int { 0 } + public var nullCount: Int { valueCount } - public init(length: Int) { - self.length = length + public init(valueCount: Int) { + self.valueCount = valueCount } public func isSet(_ bit: Int) -> Bool { - precondition(bit < length) + precondition(bit < valueCount) return false } } @@ -50,17 +56,23 @@ public struct AllNullBuffer: NullBuffer { final class BitPackedNullBuffer: NullBuffer { let length: Int let capacity: Int + let valueCount: Int let ownsMemory: Bool let buffer: UnsafePointer + var nullCount: Int init( length: Int, capacity: Int, + valueCount: Int, + nullCount: Int, ownsMemory: Bool, buffer: UnsafePointer ) { self.length = length self.capacity = capacity + self.valueCount = valueCount + self.nullCount = nullCount self.ownsMemory = ownsMemory self.buffer = buffer } diff --git a/Sources/Arrow/Buffer/NullBufferBuilder.swift b/Sources/Arrow/Buffer/NullBufferBuilder.swift index 142d13b..2414bf7 100644 --- a/Sources/Arrow/Buffer/NullBufferBuilder.swift +++ b/Sources/Arrow/Buffer/NullBufferBuilder.swift @@ -92,9 +92,9 @@ final class NullBufferBuilder { precondition(ownsMemory, "Buffer already finished.") if nullCount == 0 { - return AllValidNullBuffer(length: valueCount) + return AllValidNullBuffer(valueCount: valueCount) } else if nullCount == valueCount { - return AllNullBuffer(length: valueCount) + return AllNullBuffer(valueCount: valueCount) } else { ownsMemory = false let newCapacity = (length + 63) & ~63 @@ -107,6 +107,8 @@ final class NullBufferBuilder { return BitPackedNullBuffer( length: length, capacity: newCapacity, + valueCount: valueCount, + nullCount: nullCount, ownsMemory: true, buffer: newBuffer ) diff --git a/Sources/Arrow/Buffer/VariableLengthTypeBuffer.swift b/Sources/Arrow/Buffer/VariableLength.swift similarity index 92% rename from Sources/Arrow/Buffer/VariableLengthTypeBuffer.swift rename to Sources/Arrow/Buffer/VariableLength.swift index e8e4a93..20e1233 100644 --- a/Sources/Arrow/Buffer/VariableLengthTypeBuffer.swift +++ b/Sources/Arrow/Buffer/VariableLength.swift @@ -44,6 +44,8 @@ extension Data: VariableLength { public protocol VariableLengthBufferProtocol { associatedtype ElementType: VariableLength + var length: Int { get } + func loadVariable( at startIndex: Int, arrayLength: Int @@ -52,11 +54,11 @@ public protocol VariableLengthBufferProtocol { } /// A buffer containing values with variable length, used in variable length type Arrow arrays. -final class VariableLengthTypeBuffer: +public final class VariableLengthTypeBuffer: VariableLengthBufferProtocol { - typealias ElementType = T - var length: Int + public typealias ElementType = T + public var length: Int var capacity: Int let ownsMemory: Bool var buffer: UnsafePointer @@ -73,7 +75,7 @@ final class VariableLengthTypeBuffer: self.buffer = buffer } - func loadVariable( + public func loadVariable( at startIndex: Int, arrayLength: Int ) -> T { diff --git a/Sources/ArrowFlight/FlightClient.swift b/Sources/ArrowFlight/FlightClient.swift index 98c6a20..c305814 100644 --- a/Sources/ArrowFlight/FlightClient.swift +++ b/Sources/ArrowFlight/FlightClient.swift @@ -33,9 +33,9 @@ public class FlightClient { private func readMessages( _ responseStream: GRPCAsyncResponseStream - ) async throws -> ArrowReader.ArrowReaderResult { - let reader = ArrowReader() - let arrowResult = ArrowReader.makeArrowReaderResult() + ) async throws -> ArrowReaderX.ArrowReaderResult { + let reader = ArrowReaderX() + let arrowResult = ArrowReaderX.makeArrowReaderResult() for try await data in responseStream { switch reader.fromMessage( data.dataHeader, @@ -56,7 +56,7 @@ public class FlightClient { private func writeBatches( _ requestStream: GRPCAsyncRequestStreamWriter, descriptor: FlightDescriptor, - recordBatches: [RecordBatch] + recordBatches: [RecordBatchX] ) async throws { let writer = ArrowWriter() switch writer.toMessage(recordBatches[0].schema) { @@ -121,7 +121,7 @@ public class FlightClient { public func doGet( _ ticket: FlightTicket, - readerResultClosure: (ArrowReader.ArrowReaderResult) throws -> Void + readerResultClosure: (ArrowReaderX.ArrowReaderResult) throws -> Void ) async throws { let getResult = client.makeDoGetCall(ticket.toProtocol()) try readerResultClosure(try await readMessages(getResult.responseStream)) @@ -139,7 +139,7 @@ public class FlightClient { public func doPut( _ descriptor: FlightDescriptor, - recordBatches: [RecordBatch], + recordBatches: [RecordBatchX], closure: (FlightPutResult) throws -> Void ) async throws { if recordBatches.isEmpty { @@ -180,8 +180,8 @@ public class FlightClient { public func doExchange( _ descriptor: FlightDescriptor, - recordBatches: [RecordBatch], - closure: (ArrowReader.ArrowReaderResult) throws -> Void + recordBatches: [RecordBatchX], + closure: (ArrowReaderX.ArrowReaderResult) throws -> Void ) async throws { if recordBatches.isEmpty { throw ArrowFlightError.emptyCollection diff --git a/Sources/ArrowFlight/FlightServer.swift b/Sources/ArrowFlight/FlightServer.swift index 8367cd9..5eb9ab0 100644 --- a/Sources/ArrowFlight/FlightServer.swift +++ b/Sources/ArrowFlight/FlightServer.swift @@ -53,8 +53,8 @@ public func schemaFromMessage(_ schemaData: Data) throws(ArrowFlightError) let startIndex = schemaData.count - Int(messageLength) let schema = schemaData[startIndex...] - let reader = ArrowReader() - let result = ArrowReader.makeArrowReaderResult() + let reader = ArrowReaderX() + let result = ArrowReaderX.makeArrowReaderResult() switch reader.fromMessage(schema, dataBody: Data(), result: result) { case .success: if let resultSchema = result.schema { diff --git a/Sources/ArrowFlight/RecordBatchStreamReader.swift b/Sources/ArrowFlight/RecordBatchStreamReader.swift index 0122ea9..5d4f068 100644 --- a/Sources/ArrowFlight/RecordBatchStreamReader.swift +++ b/Sources/ArrowFlight/RecordBatchStreamReader.swift @@ -21,10 +21,10 @@ public struct RecordBatchStreamReader: AsyncSequence, AsyncIteratorProtocol, Sendable { public typealias AsyncIterator = RecordBatchStreamReader - public typealias Element = (Arrow.RecordBatch?, FlightDescriptor?) - let reader = ArrowReader() + public typealias Element = (Arrow.RecordBatchX?, FlightDescriptor?) + let reader = ArrowReaderX() // FIXME: this is hack to make this sendable - nonisolated(unsafe) var batches: [RecordBatch] = [] + nonisolated(unsafe) var batches: [RecordBatchX] = [] nonisolated(unsafe) var streamIterator: any AsyncIteratorProtocol var descriptor: FlightDescriptor? var batchIndex = 0 @@ -40,7 +40,7 @@ public struct RecordBatchStreamReader: AsyncSequence, AsyncIteratorProtocol, } public mutating func next() async throws -> ( - Arrow.RecordBatch?, FlightDescriptor? + Arrow.RecordBatchX?, FlightDescriptor? )? { guard !Task.isCancelled else { return nil @@ -52,7 +52,7 @@ public struct RecordBatchStreamReader: AsyncSequence, AsyncIteratorProtocol, return (batch, descriptor) } - let result = ArrowReader.makeArrowReaderResult() + let result = ArrowReaderX.makeArrowReaderResult() while true { let streamData = try await self.streamIterator.next() if streamData == nil { diff --git a/Sources/ArrowFlight/RecordBatchStreamWriter.swift b/Sources/ArrowFlight/RecordBatchStreamWriter.swift index cb317c8..423e722 100644 --- a/Sources/ArrowFlight/RecordBatchStreamWriter.swift +++ b/Sources/ArrowFlight/RecordBatchStreamWriter.swift @@ -57,7 +57,7 @@ public struct RecordBatchStreamWriter: Sendable { self.stream = stream } - public func write(_ rb: RecordBatch) async throws { + public func write(_ rb: RecordBatchX) async throws { // FIXME: this was moved here to make this sendable. let writer = ArrowWriter() switch writer.toMessage(rb.schema) { diff --git a/Sources/ArrowIPC/Array+IPC.swift b/Sources/ArrowIPC/Array+IPC.swift index 9eacbc9..79c5f04 100644 --- a/Sources/ArrowIPC/Array+IPC.swift +++ b/Sources/ArrowIPC/Array+IPC.swift @@ -17,7 +17,6 @@ import Foundation /// A `Data` backed Arrow utf8 array. typealias ArrowArrayUtf8 = ArrowArrayVariable< - String, FixedWidthBufferIPC, VariableLengthBufferIPC > @@ -49,7 +48,6 @@ extension ArrowArrayUtf8 { } typealias ArrowArrayBinary = ArrowArrayVariable< - Data, FixedWidthBufferIPC, VariableLengthBufferIPC > diff --git a/Sources/ArrowIPC/ArrowReader.swift b/Sources/ArrowIPC/ArrowReader.swift index a3ff2f0..eb8b347 100644 --- a/Sources/ArrowIPC/ArrowReader.swift +++ b/Sources/ArrowIPC/ArrowReader.swift @@ -31,13 +31,17 @@ struct FileDataBuffer { struct NullBufferIPC: NullBuffer { let buffer: FileDataBuffer + var valueCount: Int + let nullCount: Int + var length: Int { buffer.range.count } func isSet(_ bit: Int) -> Bool { + precondition(bit < valueCount, "Bit index \(bit) out of range") let byteIndex = bit / 8 - precondition(length > byteIndex, "Bit index \(bit) out of range") + // precondition(length > byteIndex, "Bit index \(bit) out of range") let offsetIndex = buffer.range.lowerBound + byteIndex let byte = self.buffer.data[offsetIndex] return byte & (1 << (bit % 8)) > 0 @@ -190,7 +194,7 @@ public struct ArrowReader { arrays.append(array) } - let recordBatch = RecordBatch(arrowSchema, columns: arrays) + let recordBatch = RecordBatch(schema: arrowSchema, columns: arrays) recordBatches.append(recordBatch) } @@ -199,7 +203,6 @@ public struct ArrowReader { } func loadField( - // schema: ArrowSchema, rbMessage: FRecordBatch, field: ArrowField, offset: Int64, @@ -207,10 +210,6 @@ public struct ArrowReader { bufferIndex: inout Int32 ) throws -> any ArrowArrayProtocol { - // guard let field: FField = schema.fields(at: fieldIndex) else { - // throw ArrowError.invalid("Missing field at index \(fieldIndex)") - // } - guard nodeIndex < rbMessage.nodesCount, let node = rbMessage.nodes(at: nodeIndex) else { @@ -219,10 +218,6 @@ public struct ArrowReader { nodeIndex += 1 print("incremented node index to \(nodeIndex)") - // let arrowType: ArrowType = try .type(for: field) - - let length = Int(node.length) - let buffer0 = try nextBuffer( message: rbMessage, index: &bufferIndex, @@ -231,20 +226,23 @@ public struct ArrowReader { ) // MARK: Load arrays - let nullsPresent = node.nullCount > 0 + let nullCount = Int(node.nullCount) + let length = Int(node.length) + let nullsPresent = nullCount > 0 let nullBuffer: NullBuffer if nullsPresent { - if node.nullCount == 0 { - nullBuffer = AllValidNullBuffer(length: Int(node.length)) - } else if node.length == 0 { - nullBuffer = AllValidNullBuffer(length: 0) - } else if node.nullCount == node.length { - nullBuffer = AllNullBuffer(length: Int(node.length)) + if nullCount == 0 { + nullBuffer = AllValidNullBuffer(valueCount: length) + } else if length == 0 { + nullBuffer = AllValidNullBuffer(valueCount: 0) + } else if nullCount == length { + nullBuffer = AllNullBuffer(valueCount: length) } else { - nullBuffer = NullBufferIPC(buffer: buffer0) + nullBuffer = NullBufferIPC( + buffer: buffer0, valueCount: length, nullCount: nullCount) } } else { - nullBuffer = AllValidNullBuffer(length: Int(node.length)) + nullBuffer = AllValidNullBuffer(valueCount: length) } if nullsPresent && !field.isNullable { @@ -255,14 +253,14 @@ public struct ArrowReader { if arrowType == .boolean { let buffer1 = try nextBuffer( message: rbMessage, index: &bufferIndex, offset: offset, data: data) - let valueBuffer = NullBufferIPC(buffer: buffer1) + let valueBuffer = NullBufferIPC( + buffer: buffer1, valueCount: length, nullCount: nullCount) return ArrowArrayBoolean( offset: 0, length: length, nullBuffer: nullBuffer, valueBuffer: valueBuffer) } else if arrowType.isNumeric { let buffer1 = try nextBuffer( message: rbMessage, index: &bufferIndex, offset: offset, data: data) - switch arrowType { case .float32: return makeFixedArray( @@ -275,7 +273,6 @@ public struct ArrowReader { default: throw ArrowError.notImplemented } - } else if arrowType.isVariable { let buffer1 = try nextBuffer( message: rbMessage, index: &bufferIndex, offset: offset, data: data) @@ -301,10 +298,30 @@ public struct ArrowReader { } } else if arrowType.isNested { switch arrowType { + case .list(let field): + let array: any ArrowArrayProtocol = try loadField( + rbMessage: rbMessage, + field: field, + offset: offset, + nodeIndex: &nodeIndex, + bufferIndex: &bufferIndex + ) + + let buffer1 = try nextBuffer( + message: rbMessage, index: &bufferIndex, offset: offset, data: data) + let offsetsBuffer = FixedWidthBufferIPC(buffer: buffer1) + + // This won't compile directly, so we need a helper + return makeListArray( + length: length, + nullBuffer: nullBuffer, + offsetsBuffer: offsetsBuffer, + values: array + ) + case .strct(let fields): var arrays: [(String, any ArrowArrayProtocol)] = [] for field in fields { - // print("loading field: \(field.name)") let array = try loadField( rbMessage: rbMessage, field: field, @@ -356,6 +373,21 @@ public struct ArrowReader { ) } + func makeListArray( + length: Int, + nullBuffer: NullBuffer, + offsetsBuffer: FixedWidthBufferIPC, + values: Element + ) -> AnyArrowListArray where Element: ArrowArrayProtocol { + let list = ArrowListArray( + length: length, + nullBuffer: nullBuffer, + offsetsBuffer: offsetsBuffer, + values: values + ) + return AnyArrowListArray(list) + } + private func loadSchema(_ schema: FSchema) throws(ArrowError) -> ArrowSchema { let builder = ArrowSchema.Builder() for index in 0.. Self { let type = field.typeType @@ -143,4 +143,28 @@ extension ArrowType { } } + func fType() throws(ArrowError) -> FType { + switch self { + case .int8, .int16, .int32, .int64, .uint8, .uint16, .uint32, .uint64: + return .int + case .float16, .float32, .float64: + return .floatingpoint + case .utf8: + return .utf8 + case .binary: + return .binary + case .boolean: + return .bool + case .date32, .date64: + return .date + case .time32, .time64: + return .time + case .timestamp: + return .timestamp + case .strct: + return .struct_ + default: + throw .invalid("Unhandled field type: \(self)") + } + } } diff --git a/Sources/ArrowIPC/ArrowWriter+recordBatch.swift b/Sources/ArrowIPC/ArrowWriter+recordBatch.swift new file mode 100644 index 0000000..dc76270 --- /dev/null +++ b/Sources/ArrowIPC/ArrowWriter+recordBatch.swift @@ -0,0 +1,133 @@ +// Copyright 2025 The Apache Software Foundation +// Copyright 2025 The Columnar Swift Contributors +// +// Licensed 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 Arrow +import FlatBuffers +import Foundation + +extension ArrowWriter { + + private func write( + batch: RecordBatch + ) throws -> Offset { + let schema = batch.schema + var fbb = FlatBufferBuilder() + + // write out field nodes + var fieldNodeOffsets: [Offset] = [] + fbb.startVector( + schema.fields.count, + elementSize: MemoryLayout.size + ) + writeFieldNodes( + fields: schema.fields, + columns: batch.columns, + offsets: &fieldNodeOffsets, + fbb: &fbb + ) + let nodeOffset = fbb.endVector(len: fieldNodeOffsets.count) + // write out buffers + var buffers: [FBuffer] = .init() + var bufferOffset: Int = 0 + writeBufferInfo( + schema.fields, columns: batch.columns, + bufferOffset: &bufferOffset, buffers: &buffers, + fbb: &fbb + ) + FRecordBatch.startVectorOfBuffers(batch.schema.fields.count, in: &fbb) + for buffer in buffers.reversed() { + fbb.create(struct: buffer) + } + let batchBuffersOffset = fbb.endVector(len: buffers.count) + let startRb = FRecordBatch.startRecordBatch(&fbb) + FRecordBatch.addVectorOf(nodes: nodeOffset, &fbb) + FRecordBatch.addVectorOf(buffers: batchBuffersOffset, &fbb) + FRecordBatch.add(length: Int64(batch.length), &fbb) + let recordBatchOffset = FRecordBatch.endRecordBatch( + &fbb, + start: startRb + ) + let bodySize = Int64(bufferOffset) + let startMessage = FMessage.startMessage(&fbb) + FMessage.add(version: .max, &fbb) + FMessage.add(bodyLength: Int64(bodySize), &fbb) + FMessage.add(headerType: .recordbatch, &fbb) + FMessage.add(header: recordBatchOffset, &fbb) + let messageOffset = FMessage.endMessage(&fbb, start: startMessage) + fbb.finish(offset: messageOffset) + // return .success((fbb.data, Offset(offset: UInt32(fbb.data.count)))) + fatalError() + } + + private func writeFieldNodes( + fields: [ArrowField], + columns: [any ArrowArrayProtocol], + offsets: inout [Offset], + fbb: inout FlatBufferBuilder + ) { + fatalError() + // for index in (0.. Int { + let padding = byteCount % alignment + if padding > 0 { + return byteCount + alignment - padding + } + return byteCount + } + + mutating func write(bytes: [UInt8], alignment: Int) { + data.append(contentsOf: fileMarker) + let remainder = bytes.count % alignment + let padding = alignment - remainder + if padding > 0 { + data.append(contentsOf: [UInt8](repeating: 0, count: padding)) + } + precondition(data.count % 8 == 0, "File must be aligned to 8 bytes.") + } + + mutating func write(data other: Data, alignment: Int = 8) { + self.data.append(other) + let remainder = data.count % alignment + let padding = alignment - remainder + if padding > 0 { + data.append(contentsOf: [UInt8](repeating: 0, count: padding)) + } + precondition(data.count % 8 == 0, "File must be aligned to 8 bytes.") + } + + mutating func write(recordBatches: [RecordBatch]) throws { + + } + + mutating func write(schema: ArrowSchema) throws(ArrowError) { + var fbb: FlatBufferBuilder = .init() + let schemaOffset = try write(schema: schema, to: &fbb) + fbb.finish(offset: schemaOffset) + self.write(data: fbb.data) + } + + private func write( + schema: ArrowSchema, + to fbb: inout FlatBufferBuilder + ) throws(ArrowError) -> Offset { + var fieldOffsets: [Offset] = [] + for field in schema.fields { + let offset = try write(field: field, to: &fbb) + fieldOffsets.append(offset) + } + let fieldsOffset: Offset = fbb.createVector(ofOffsets: fieldOffsets) + let schemaOffset = FSchema.createSchema( + &fbb, + endianness: .little, + fieldsVectorOffset: fieldsOffset + ) + return schemaOffset + } + + private func write( + field: ArrowField, + to fbb: inout FlatBufferBuilder, + ) throws(ArrowError) -> Offset { + var fieldsOffset: Offset? + if case .strct(let fields) = field.type { + var offsets: [Offset] = [] + for field in fields { + let offset = try write(field: field, to: &fbb) + offsets.append(offset) + } + fieldsOffset = fbb.createVector(ofOffsets: offsets) + } + let nameOffset = fbb.create(string: field.name) + let fieldTypeOffset = try append(arrowType: field.type, to: &fbb) + let startOffset = FField.startField(&fbb) + FField.add(name: nameOffset, &fbb) + FField.add(nullable: field.isNullable, &fbb) + if let childrenOffset = fieldsOffset { + FField.addVectorOf(children: childrenOffset, &fbb) + } + let typeType = try field.type.fType() + FField.add(typeType: typeType, &fbb) + FField.add(type: fieldTypeOffset, &fbb) + return FField.endField(&fbb, start: startOffset) + } + + private func append( + arrowType: ArrowType, + to fbb: inout FlatBufferBuilder, + ) throws(ArrowError) -> Offset { + switch arrowType { + case .int8, .uint8: + return FInt.createInt(&fbb, bitWidth: 8, isSigned: arrowType == .int8) + case .int16, .uint16: + return FInt.createInt(&fbb, bitWidth: 16, isSigned: arrowType == .int16) + case .int32, .uint32: + return FInt.createInt(&fbb, bitWidth: 32, isSigned: arrowType == .int32) + case .int64, .uint64: + return FInt.createInt(&fbb, bitWidth: 64, isSigned: arrowType == .int64) + case .float16: + return FFloatingPoint.createFloatingPoint(&fbb, precision: .half) + case .float32: + return FFloatingPoint.createFloatingPoint(&fbb, precision: .single) + case .float64: + return FFloatingPoint.createFloatingPoint(&fbb, precision: .double) + case .utf8: + return FUtf8.endUtf8(&fbb, start: FUtf8.startUtf8(&fbb)) + case .binary: + return FBinary.endBinary(&fbb, start: FBinary.startBinary(&fbb)) + case .boolean: + return FBool.endBool(&fbb, start: FBool.startBool(&fbb)) + case .date32: + let startOffset = FDate.startDate(&fbb) + FDate.add(unit: .day, &fbb) + return FDate.endDate(&fbb, start: startOffset) + case .date64: + let startOffset = FDate.startDate(&fbb) + FDate.add(unit: .millisecond, &fbb) + return FDate.endDate(&fbb, start: startOffset) + case .time32(let unit): + let startOffset = FTime.startTime(&fbb) + FTime.add(unit: unit == .second ? .second : .millisecond, &fbb) + return FTime.endTime(&fbb, start: startOffset) + case .time64(let unit): + let startOffset = FTime.startTime(&fbb) + FTime.add(unit: unit == .microsecond ? .microsecond : .nanosecond, &fbb) + return FTime.endTime(&fbb, start: startOffset) + case .timestamp(let unit, let timezone): + let startOffset = FTimestamp.startTimestamp(&fbb) + let fbUnit: FTimeUnit + switch unit { + case .second: + fbUnit = .second + case .millisecond: + fbUnit = .millisecond + case .microsecond: + fbUnit = .microsecond + case .nanosecond: + fbUnit = .nanosecond + } + FTimestamp.add(unit: fbUnit, &fbb) + if let timezone { + let timezoneOffset = fbb.create(string: timezone) + FTimestamp.add(timezone: timezoneOffset, &fbb) + } + return FTimestamp.endTimestamp(&fbb, start: startOffset) + case .strct(_): + let startOffset = FStruct.startStruct_(&fbb) + return FStruct.endStruct_(&fbb, start: startOffset) + default: + throw .unknownType( + "Unable to add flatbuf type for Arrow type: \(arrowType)") + } + } +} diff --git a/Sources/ArrowIPC/RecordBatch.swift b/Sources/ArrowIPC/RecordBatch.swift index f22d5e4..2af6e6e 100644 --- a/Sources/ArrowIPC/RecordBatch.swift +++ b/Sources/ArrowIPC/RecordBatch.swift @@ -1,27 +1,27 @@ -// RecordBatch.swift -// Arrow +// Copyright 2025 The Columnar Swift Contributors // -// Created by Will Temperley on 14/11/2025. All rights reserved. -// Copyright 2025 Will Temperley. +// Licensed 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 // -// Copying or reproduction of this file via any medium requires prior express -// written permission from the copyright holder. -// ----------------------------------------------------------------------------- -/// -/// Implementation notes, links and internal documentation go here. -/// -// ----------------------------------------------------------------------------- +// 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 Arrow -struct RecordBatch { +public struct RecordBatch { let schema: ArrowSchema - var columnCount: Int { Int(self.columns.count) } + var columnCount: Int { self.columns.count } let columns: [any ArrowArrayProtocol] let length: Int - public init(_ schema: ArrowSchema, columns: [any ArrowArrayProtocol]) { + public init(schema: ArrowSchema, columns: [any ArrowArrayProtocol]) { self.schema = schema self.columns = columns self.length = columns[0].length diff --git a/Tests/ArrowFlightTests/FlightTest.swift b/Tests/ArrowFlightTests/FlightTest.swift index 7a0bed4..1054e15 100644 --- a/Tests/ArrowFlightTests/FlightTest.swift +++ b/Tests/ArrowFlightTests/FlightTest.swift @@ -34,7 +34,7 @@ func makeSchema() -> ArrowSchema { .finish() } -func makeRecordBatch() throws -> RecordBatch { +func makeRecordBatch() throws -> RecordBatchX { let doubleBuilder: NumberArrayBuilder = try ArrowArrayBuilders.loadNumberArrayBuilder() doubleBuilder.append(11.11) @@ -56,7 +56,7 @@ func makeRecordBatch() throws -> RecordBatch { let doubleArray = try doubleBuilder.finish() let stringArray = try stringBuilder.finish() let date32Array = try date32Builder.finish() - let result = RecordBatch.Builder() + let result = RecordBatchX.Builder() .addColumn("col1", arrowArray: doubleArray) .addColumn("col2", arrowArray: stringArray) .addColumn("col3", arrowArray: date32Array) @@ -295,8 +295,8 @@ public class FlightClientTester { func doGetTestFlightData(_ flightData: Data) async throws { let ticket = FlightTicket(flightData) var numCall = 0 - let reader = ArrowReader() - let arrowResult = ArrowReader.makeArrowReaderResult() + let reader = ArrowReaderX() + let arrowResult = ArrowReaderX.makeArrowReaderResult() try await client?.doGet( ticket, flightDataClosure: { flightData in diff --git a/Tests/ArrowIPCTests/ArrowWriterTests.swift b/Tests/ArrowIPCTests/ArrowWriterTests.swift new file mode 100644 index 0000000..e3623dd --- /dev/null +++ b/Tests/ArrowIPCTests/ArrowWriterTests.swift @@ -0,0 +1,109 @@ +// Copyright 2025 The Columnar Swift Contributors +// +// Licensed 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 +import Testing + +@testable import Arrow +@testable import ArrowIPC + +struct ArrowWriterTests { + + @Test func writeBasics() throws { + + let outputUrl = FileManager.default.temporaryDirectory + .appending(path: "bool-test.arrow") + let writer = ArrowWriter(url: outputUrl) + #expect(writer.data.count == 8) + + } + + @Test func writeBoolean() throws { + + let schema: ArrowSchema = ArrowSchema.Builder() + .addField("one", type: .boolean, isNullable: true) + .addField("two", type: .utf8, isNullable: true) + .finish() + + let builder = ArrayBuilderBoolean() + builder.append(true) + builder.append(false) + builder.appendNull() + builder.append(false) + builder.append(true) + let one = builder.finish() + + let builder2 = ArrayBuilderString() + builder2.append("zero") + builder2.append("one") + builder2.append("two") + builder2.append("three") + builder2.append("four") + let two = builder2.finish() + + let recordBatch = RecordBatch(schema: schema, columns: [one, two]) + + let outputUrl = FileManager.default.temporaryDirectory + .appending(path: "bool-test.arrow") + let writer = ArrowWriter(url: outputUrl) + + // writer.write(recordBatch: recordBatch) + + // func writeBoolData() { + // alloc := memory.NewGoAllocator() + // schema := arrow.NewSchema([]arrow.Field{ + // {Name: "one", Type: arrow.FixedWidthTypes.Boolean}, + // {Name: "two", Type: arrow.BinaryTypes.String}, + // }, nil) + // + // b := array.NewRecordBuilder(alloc, schema) + // defer b.Release() + // + // b.Field(0).(*array.BooleanBuilder).AppendValues([]bool{true, false}, nil) + // b.Field(0).(*array.BooleanBuilder).AppendNull() + // b.Field(0).(*array.BooleanBuilder).AppendValues([]bool{false, true}, nil) + // b.Field(1).(*array.StringBuilder).AppendValues([]string{"zero", "one", "two", "three", "four"}, nil) + // rec := b.NewRecord() + // defer rec.Release() + // + // writeBytes(rec, "testdata_bool.arrow") + // } + + // // read existing file + // let fileURL = try loadArrowResource(name: "testdata_bool") + // let arrowReader = ArrowReader() + //// let fileRBs = try checkBoolRecordBatch(arrowReader.fromFile(fileURL)) + // let arrowWriter = ArrowWriter() + // // write data from file to a stream + // let writerInfo = ArrowWriter.Info( + // .recordbatch, schema: fileRBs[0].schema, batches: fileRBs) + // switch arrowWriter.writeFile(writerInfo) { + // case .success(let writeData): + // // read stream back into recordbatches + // try checkBoolRecordBatch(arrowReader.readFile(writeData)) + // case .failure(let error): + // throw error + // } + // // write file record batches to another file + // let outputUrl = FileManager.default.temporaryDirectory + // .appending(path: "testfilewriter_bool.arrow") + // switch arrowWriter.toFile(outputUrl, info: writerInfo) { + // case .success: + // try checkBoolRecordBatch(arrowReader.fromFile(outputUrl)) + // case .failure(let error): + // throw error + // } + } + +} diff --git a/Tests/ArrowTests/Array/ArrayNullBufferTests.swift b/Tests/ArrowTests/Array/ArrayNullBufferTests.swift new file mode 100644 index 0000000..a75df24 --- /dev/null +++ b/Tests/ArrowTests/Array/ArrayNullBufferTests.swift @@ -0,0 +1,50 @@ +// ArrayNullBufferTests.swift +// Arrow +// +// Created by Will Temperley on 19/11/2025. All rights reserved. +// Copyright 2025 Will Temperley. +// +// Copying or reproduction of this file via any medium requires prior express +// written permission from the copyright holder. +// ----------------------------------------------------------------------------- +/// +/// Implementation notes, links and internal documentation go here. +/// +// ----------------------------------------------------------------------------- + +import Testing + +@testable import Arrow + +/// Test internal null buffer types are correct. +struct ArrayNullBufferTests { + + @Test func allValidValues() throws { + // Should be able to omit null buffer entirely + let arrayBuilder: ArrayBuilderFixedWidth = .init() + for i in 0..<1000 { + arrayBuilder.append(Int64(i)) // No nulls + } + let array = arrayBuilder.finish() + for i in 0..<1000 { + #expect(array[i]! == Int64(i)) + } + let nullBuffer = try #require(array.nullBuffer as? AllValidNullBuffer) + #expect(nullBuffer.valueCount == 1000) + #expect(array.bufferSizes == [0, 1000 * MemoryLayout.stride]) + } + + @Test func allNullValues() throws { + let arrayBuilder: ArrayBuilderFixedWidth = .init() + for _ in 0..<1000 { + arrayBuilder.appendNull() + } + let array = arrayBuilder.finish() + for i in 0..<1000 { + #expect(array[i] == nil) + } + let nullBuffer = try #require(array.nullBuffer as? AllNullBuffer) + #expect(nullBuffer.valueCount == 1000) + #expect(array.bufferSizes == [0, 1000 * MemoryLayout.stride]) + } +} diff --git a/Tests/ArrowTests/Array/ArrayTests2.swift b/Tests/ArrowTests/Array/ArrayTests2.swift deleted file mode 100644 index ae6ea44..0000000 --- a/Tests/ArrowTests/Array/ArrayTests2.swift +++ /dev/null @@ -1,751 +0,0 @@ -// Copyright 2025 The Apache Software Foundation -// Copyright 2025 The Columnar Swift Contributors -// -// Licensed 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 -import Testing - -@testable import Arrow - -struct ArrayTests2 { - - @Test func allValidValues() throws { - // Should be able to omit null buffer entirely - let arrayBuilder: ArrayBuilderFixedWidth = .init() - for i in 0..<1000 { - arrayBuilder.append(Int64(i)) // No nulls - } - let array = arrayBuilder.finish() - - #expect(array.nullBuffer is AllValidNullBuffer) - for i in 0..<1000 { - #expect(array[i]! == Int64(i)) - } - } - - @Test func allNullValues() throws { - // Must have null buffer, all bits = 0 - let arrayBuilder: ArrayBuilderFixedWidth = .init() - for _ in 0..<1000 { - arrayBuilder.appendNull() - } - let array = arrayBuilder.finish() - #expect(array.nullBuffer is AllNullBuffer) - for i in 0..<1000 { - #expect(array[i] == nil) - } - // Null buffer must exist, all nulls - } - - @Test func boolArray() throws { - - let builder = ArrayBuilderBoolean() - builder.append(true) - builder.appendNull() - builder.append(false) - builder.append(false) - let boolArray = builder.finish() - #expect(boolArray.length == 4) - #expect(boolArray[0]! == true) - #expect(boolArray[1] == nil) - #expect(boolArray[2]! == false) - #expect(boolArray[3]! == false) - } - - @Test func uint8Array() throws { - - // MARK: UInt8 array - let arrayBuilder: ArrayBuilderFixedWidth = .init() - for index: UInt8 in 0..<100 { - arrayBuilder.append(index) - } - - arrayBuilder.appendNull() - #expect(arrayBuilder.length == 101) - let array = arrayBuilder.finish() - #expect(array.length == 101) - #expect(array[1]! == 1) - #expect(array[10]! == 10) - #expect(array[100] == nil) - - for index in 0..<100 { - #expect(array[Int(index)]! == UInt8(index)) - } - - let slice = array.slice(offset: 5, length: 5) - for i in 0..<5 { - #expect(slice[i] == UInt8(5 + i)) - } - } - - @Test func int64Array() throws { - var rng = getSeededRNG() - let count = Int.random(in: 0...100_000) - var testArray = [Int64](repeating: 0, count: count) - for i in 0.. = .init() - for i in 0.. = .init() - - builder.appendNull() - builder.append("abc") - builder.append("def") - builder.appendNull() - builder.append("This is a longer string") - builder.appendNull() - builder.appendNull() - for i in 0..<100 { - builder.append("test \(i)") - } - - let stringArray = builder.finish() - #expect(stringArray[0] == nil) - #expect(stringArray[1]! == "abc") - #expect(stringArray[2]! == "def") - #expect(stringArray[3] == nil) - #expect(stringArray[4]! == "This is a longer string") - #expect(stringArray[5] == nil) - #expect(stringArray[6] == nil) - for i in 0..<100 { - #expect(stringArray[i + 7]! == "test \(i)") - } - } - - @Test func stringArrayWithRandomNulls() throws { - var rng = getSeededRNG() - let count = Int.random(in: 0...100_000) - var testArray = [String?](repeating: nil, count: count) - - // Generate random strings with random nulls - for i in 0.. = .init() - for value in testArray { - if let value { - arrayBuilder.append(value) - } else { - arrayBuilder.appendNull() - } - } - let stringArray = arrayBuilder.finish() - - for i in 0.. = .init() - for index in 0..<100 { - if index % 10 == 9 { - arrayBuilder.appendNull() - } else { - let val = Data("test\(index)".utf8) - arrayBuilder.append(val) - } - } - - let binaryArray = arrayBuilder.finish() - #expect(binaryArray.length == 100) - for index in 0.. = .init() - for value in expected { - if let value { - arrayBuilder.append(value) - } else { - arrayBuilder.appendNull() - } - } - let binaryArray = arrayBuilder.finish() - - for i in 0.. = .init() - for value in expected { - if let value { - arrayBuilder.append(value) - } else { - arrayBuilder.appendNull() - } - } - let int64Array = arrayBuilder.finish() - - for i in 0.. nullProbability { - let length = Int.random(in: 0...50, using: &rng) - expected[i] = randomString(length: length, using: &rng) - } - } - let arrayBuilder: ArrayBuilderVariable = .init() - for value in expected { - if let value { - arrayBuilder.append(value) - } else { - arrayBuilder.appendNull() - } - } - let stringArray = arrayBuilder.finish() - - for i in 0.. = .init() - for value in expected { - if let value { - arrayBuilder.append(value) - } else { - arrayBuilder.appendNull() - } - } - let stringArray = arrayBuilder.finish() - - for i in 0.. = .init() - for value in expected { - if let value { - arrayBuilder.append(value) - } else { - arrayBuilder.appendNull() - } - } - let int64Array = arrayBuilder.finish() - for i in 0.. = .init() - builder.append(14) - builder.appendNull() - builder.append(40.4) - let doubleArray = builder.finish() - #expect(doubleArray.length == 3) - #expect(doubleArray[0]! == 14) - #expect(doubleArray[1] == nil) - #expect(doubleArray[2]! == 40.4) - } - - @Test func date32Array() throws { - let date32Builder: ArrayBuilderDate32 = .init() - let date2 = Date(timeIntervalSinceReferenceDate: 86400 * 1) - let date1 = Date(timeIntervalSinceReferenceDate: 86400 * 5000 + 352) - date32Builder.append(date1) - date32Builder.append(date2) - date32Builder.appendNull() - let date32Array = date32Builder.finish() - #expect(date32Array.length == 3) - #expect(date32Array[1] == date2) - let adjustedDate1 = Date( - timeIntervalSince1970: date1.timeIntervalSince1970 - 352) - #expect(date32Array[0]! == adjustedDate1) - } - - @Test func date64Array() throws { - let date64Builder: ArrayBuilderDate64 = .init() - let date2 = Date(timeIntervalSinceReferenceDate: 86400 * 1) - let date1 = Date(timeIntervalSinceReferenceDate: 86400 * 5000 + 352) - date64Builder.append(date1) - date64Builder.append(date2) - date64Builder.appendNull() - let date64Array = date64Builder.finish() - #expect(date64Array.length == 3) - #expect(date64Array[1] == date2) - #expect(date64Array[0]! == date1) - } - - @Test func time32Array() throws { - let milliBuilder: ArrayBuilderTime32 = .init() - milliBuilder.append(100) - milliBuilder.append(1_000_000) - milliBuilder.appendNull() - let milliArray = milliBuilder.finish() - #expect(milliArray.length == 3) - #expect(milliArray[1] == 1_000_000) - #expect(milliArray[2] == nil) - - let secBuilder: ArrayBuilderTime32 = .init() - secBuilder.append(200) - secBuilder.appendNull() - secBuilder.append(2_000_011) - let secArray = secBuilder.finish() - #expect(secArray.length == 3) - #expect(secArray[1] == nil) - #expect(secArray[2] == 2_000_011) - } - - @Test func time64Array() throws { - let nanoBuilder: ArrayBuilderTime64 = .init() - nanoBuilder.append(10000) - nanoBuilder.appendNull() - nanoBuilder.append(123_456_789) - let nanoArray = nanoBuilder.finish() - #expect(nanoArray.length == 3) - #expect(nanoArray[1] == nil) - #expect(nanoArray[2] == 123_456_789) - - let microBuilder: ArrayBuilderTime64 = .init() - microBuilder.appendNull() - microBuilder.append(20000) - microBuilder.append(987_654_321) - - let microArray = microBuilder.finish() - #expect(microArray.length == 3) - #expect(microArray[1] == 20000) - #expect(microArray[2] == 987_654_321) - } - - @Test func timestampArray() throws { - // Test timestamp with seconds unit - let secBuilder: ArrayBuilderTimestamp = .init() - secBuilder.append(1_609_459_200) // 2021-01-01 00:00:00 - secBuilder.append(1_609_545_600) // 2021-01-02 00:00:00 - secBuilder.appendNull() - let secArray = secBuilder.finish() - #expect(secArray.length == 3) - #expect(secArray[0] == 1_609_459_200) - #expect(secArray[1] == 1_609_545_600) - #expect(secArray[2] == nil) - - // Test timestamp with milliseconds unit and timezone America/New_York - let msBuilder: ArrayBuilderTimestamp = .init() - msBuilder.append(1_609_459_200_000) // 2021-01-01 00:00:00.000 - msBuilder.appendNull() - msBuilder.append(1_609_545_600_000) // 2021-01-02 00:00:00.000 - let msArray = msBuilder.finish() - #expect(msArray.length == 3) - #expect(msArray[0] == 1_609_459_200_000) - #expect(msArray[1] == nil) - #expect(msArray[2] == 1_609_545_600_000) - - // Test timestamp with microseconds unit and timezone UTC - let usBuilder: ArrayBuilderTimestamp = .init() - usBuilder.append(1_609_459_200_000_000) // 2021-01-01 00:00:00.000000 - usBuilder.append(1_609_545_600_000_000) // 2021-01-02 00:00:00.000000 - usBuilder.append(1_609_632_000_000_000) // 2021-01-03 00:00:00.000000 - let usArray = usBuilder.finish() - #expect(usArray.length == 3) - #expect(usArray[0] == 1_609_459_200_000_000) - #expect(usArray[1] == 1_609_545_600_000_000) - #expect(usArray[2] == 1_609_632_000_000_000) - - // Test timestamp with nanoseconds unit - let nsBuilder: ArrayBuilderTimestamp = .init() - nsBuilder.appendNull() - // 2021-01-01 00:00:00.000000000 - nsBuilder.append(1_609_459_200_000_000_000) - // 2021-01-02 00:00:00.000000000 - nsBuilder.append(1_609_545_600_000_000_000) - let nsArray = nsBuilder.finish() - #expect(nsArray.length == 3) - #expect(nsArray[0] == nil) - #expect(nsArray[1] == 1_609_459_200_000_000_000) - #expect(nsArray[2] == 1_609_545_600_000_000_000) - } - - // MARK: need to migrate these - @Test func structArray() throws { - class StructTest { - var fieldBool: Bool = false - var fieldInt8: Int8 = 0 - var fieldInt16: Int16 = 0 - var fieldInt32: Int32 = 0 - var fieldInt64: Int64 = 0 - var fieldUInt8: UInt8 = 0 - var fieldUInt16: UInt16 = 0 - var fieldUInt32: UInt32 = 0 - var fieldUInt64: UInt64 = 0 - var fieldDouble: Double = 0 - var fieldFloat: Float = 0 - var fieldString: String = "" - var fieldData = Data() - var fieldDate: Date = Date.now - } - - enum STIndex: Int { - case bool, int8, int16, int32, int64 - case uint8, uint16, uint32, uint64, double - case float, string, data, date - } - - let testData = StructTest() - let dateNow = Date.now - let structBuilder = try ArrowArrayBuilders.structArrayBuilderForType( - testData) - structBuilder.append([ - true, Int8(1), Int16(2), Int32(3), Int64(4), - UInt8(5), UInt16(6), UInt32(7), UInt64(8), Double(9.9), - Float(10.10), "11", Data("12".utf8), dateNow, - ]) - structBuilder.append(nil) - structBuilder.append([ - true, Int8(13), Int16(14), Int32(15), Int64(16), - UInt8(17), UInt16(18), UInt32(19), UInt64(20), Double(21.21), - Float(22.22), "23", Data("24".utf8), dateNow, - ]) - #expect(structBuilder.length == 3) - let structArray = try structBuilder.finish() - #expect(structArray.length == 3) - #expect(structArray[1] == nil) - - #expect(structArray.fields![0].length == 3) - #expect(structArray.fields![0].asAny(1) == nil) - #expect(structArray[0]![STIndex.bool.rawValue] as? Bool == true) - #expect(structArray[0]![STIndex.int8.rawValue] as? Int8 == 1) - #expect(structArray[0]![STIndex.int16.rawValue] as? Int16 == 2) - #expect(structArray[0]![STIndex.int32.rawValue] as? Int32 == 3) - #expect(structArray[0]![STIndex.int64.rawValue] as? Int64 == 4) - #expect(structArray[0]![STIndex.uint8.rawValue] as? UInt8 == 5) - #expect(structArray[0]![STIndex.uint16.rawValue] as? UInt16 == 6) - #expect(structArray[0]![STIndex.uint32.rawValue] as? UInt32 == 7) - #expect(structArray[0]![STIndex.uint64.rawValue] as? UInt64 == 8) - #expect(structArray[0]![STIndex.double.rawValue] as? Double == 9.9) - #expect(structArray[0]![STIndex.float.rawValue] as? Float == 10.10) - #expect(structArray[2]![STIndex.string.rawValue] as? String == "23") - #expect( - String( - decoding: (structArray[0]![STIndex.data.rawValue] as? Data)!, - as: UTF8.self) == "12") - let dateFormatter = DateFormatter() - dateFormatter.timeStyle = .full - #expect( - dateFormatter.string( - from: (structArray[0]![STIndex.date.rawValue] as? Date)!) - == dateFormatter.string(from: dateNow)) - } - - func checkHolderForType(_ checkType: ArrowType) throws { - let buffers = [ - ArrowBuffer( - length: 0, capacity: 0, - rawPointer: UnsafeMutableRawPointer.allocate( - byteCount: 0, alignment: .zero)), - ArrowBuffer( - length: 0, capacity: 0, - rawPointer: UnsafeMutableRawPointer.allocate( - byteCount: 0, alignment: .zero)), - ] - let field = ArrowField(name: "", dataType: checkType, isNullable: true) - switch makeArrayHolder( - field, buffers: buffers, nullCount: 0, children: nil, rbLength: 0) - { - case .success(let holder): - #expect(holder.type == checkType) - case .failure(let err): - throw err - } - } - - @Test func arrayHolders() throws { - try checkHolderForType(.int8) - try checkHolderForType(.uint8) - try checkHolderForType(.int16) - try checkHolderForType(.uint16) - try checkHolderForType(.int32) - try checkHolderForType(.uint32) - try checkHolderForType(.int64) - try checkHolderForType(.uint64) - try checkHolderForType(.time32(.second)) - try checkHolderForType(.time32(.millisecond)) - try checkHolderForType(.time64(.microsecond)) - try checkHolderForType(.time64(.nanosecond)) - try checkHolderForType(.binary) - try checkHolderForType(.float32) - try checkHolderForType(.float64) - try checkHolderForType(.boolean) - try checkHolderForType(.utf8) - } - - @Test func arrowArrayHolderBuilder() throws { - let uint8HBuilder: AnyArrowArrayBuilder = - (try ArrowArrayBuilders.loadNumberArrayBuilder() - as NumberArrayBuilder) - for index in 0..<100 { - uint8HBuilder.appendAny(UInt8(index)) - } - - let uint8Holder = try uint8HBuilder.toAnyArrowArray() - #expect(uint8Holder.nullCount == 0) - #expect(uint8Holder.length == 100) - - let stringHBuilder: AnyArrowArrayBuilder = - (try ArrowArrayBuilders.loadStringArrayBuilder()) - for index in 0..<100 { - if index % 10 == 9 { - stringHBuilder.appendAny(nil) - } else { - stringHBuilder.appendAny("test" + String(index)) - } - } - - let stringHolder = try stringHBuilder.toAnyArrowArray() - #expect(stringHolder.nullCount == 10) - #expect(stringHolder.length == 100) - } - - @Test func addVArgs() throws { - let arrayBuilder: NumberArrayBuilder = - try ArrowArrayBuilders.loadNumberArrayBuilder() - arrayBuilder.append(0, 1, 2, 3, 4, 5, 6, 7, 8, 9) - #expect(arrayBuilder.length == 10) - #expect(try arrayBuilder.finish()[2] == 2) - let doubleBuilder: NumberArrayBuilder = - try ArrowArrayBuilders.loadNumberArrayBuilder() - doubleBuilder.append(0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8) - #expect(doubleBuilder.length == 9) - #expect(try doubleBuilder.finish()[4] == 4.4) - let stringBuilder = try ArrowArrayBuilders.loadStringArrayBuilder() - stringBuilder.append("0", "1", "2", "3", "4", "5", "6") - #expect(stringBuilder.length == 7) - #expect(try stringBuilder.finish()[4] == "4") - let boolBuilder = try ArrowArrayBuilders.loadBoolArrayBuilder() - boolBuilder.append(true, false, true, false) - #expect(try boolBuilder.finish()[2] == true) - } - - @Test func addArray() throws { - let arrayBuilder: NumberArrayBuilder = - try ArrowArrayBuilders.loadNumberArrayBuilder() - arrayBuilder.append([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) - #expect(arrayBuilder.length == 10) - #expect(try arrayBuilder.finish()[2] == 2) - let doubleBuilder: NumberArrayBuilder = - try ArrowArrayBuilders.loadNumberArrayBuilder() - doubleBuilder.append([0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8]) - #expect(doubleBuilder.length == 9) - #expect(try doubleBuilder.finish()[4] == 4.4) - let stringBuilder = try ArrowArrayBuilders.loadStringArrayBuilder() - stringBuilder.append(["0", "1", "2", "3", "4", "5", "6"]) - #expect(stringBuilder.length == 7) - #expect(try stringBuilder.finish()[4] == "4") - let boolBuilder = try ArrowArrayBuilders.loadBoolArrayBuilder() - boolBuilder.append([true, false, true, false]) - #expect(try boolBuilder.finish()[2] == true) - } - - @Test func listArrayPrimitive() throws { - let field = ArrowField(listFieldWith: .int32, isNullable: false) - let listBuilder = try ListArrayBuilder(.list(field)) - - listBuilder.append([Int32(1), Int32(2), Int32(3)]) - listBuilder.append([Int32(4), Int32(5)]) - listBuilder.append(nil) - listBuilder.append([Int32(6), Int32(7), Int32(8), Int32(9)]) - - #expect(listBuilder.length == 4) - #expect(listBuilder.nullCount == 1) - - let listArray = try listBuilder.finish() - #expect(listArray.length == 4) - - let firstList = listArray[0] - #expect(firstList != nil, "First list should not be nil") - #expect(firstList!.count == 3, "First list should have 3 elements") - #expect(firstList![0] as? Int32 == 1) - #expect(firstList![1] as? Int32 == 2) - #expect(firstList![2] as? Int32 == 3) - - let secondList = listArray[1] - #expect(secondList!.count == 2) - #expect(secondList![0] as? Int32 == 4) - #expect(secondList![1] as? Int32 == 5) - - #expect(listArray[2] == nil) - - let fourthList = listArray[3] - #expect(fourthList!.count == 4) - #expect(fourthList![0] as? Int32 == 6) - #expect(fourthList![3] as? Int32 == 9) - } - - @Test func listArrayNested() throws { - let field = ArrowField(listFieldWith: .int32, isNullable: false) - let innerListType: ArrowType = .list(field) - let outerField = ArrowField(listFieldWith: innerListType, isNullable: false) - let outerListBuilder = try ListArrayBuilder(.list(outerField)) - - guard - let innerListBuilder = outerListBuilder.valueBuilder as? ListArrayBuilder - else { - Issue.record("Failed to cast valueBuilder to ListArrayBuilder") - return - } - - outerListBuilder.bufferBuilder.append(2) - innerListBuilder.append([Int32(1), Int32(2)]) - innerListBuilder.append([Int32(3), Int32(4), Int32(5)]) - - outerListBuilder.bufferBuilder.append(1) - innerListBuilder.append([Int32(6)]) - - outerListBuilder.bufferBuilder.append(nil) - - outerListBuilder.bufferBuilder.append([]) - - let nestedArray = try outerListBuilder.finish() - #expect(nestedArray.length == 4) - #expect(nestedArray.nullCount == 1) - - let firstOuterList = nestedArray[0]! - #expect(firstOuterList.count == 2) - - let firstInnerList = firstOuterList[0] as! [Any?] - #expect(firstInnerList.count == 2) - #expect(firstInnerList[0] as? Int32 == 1) - #expect(firstInnerList[1] as? Int32 == 2) - - let secondInnerList = firstOuterList[1] as! [Any?] - #expect(secondInnerList.count == 3) - #expect(secondInnerList[0] as? Int32 == 3) - #expect(secondInnerList[1] as? Int32 == 4) - #expect(secondInnerList[2] as? Int32 == 5) - - let secondOuterList = nestedArray[1]! - #expect(secondOuterList.count == 1) - - let thirdInnerList = secondOuterList[0] as! [Any?] - #expect(thirdInnerList.count == 1) - #expect(thirdInnerList[0] as? Int32 == 6) - - #expect(nestedArray[2] == nil) - - let emptyList = nestedArray[3]! - #expect(emptyList.count == 0) - } -} diff --git a/Tests/ArrowTests/Array/BasicArrayTests.swift b/Tests/ArrowTests/Array/BasicArrayTests.swift new file mode 100644 index 0000000..59bd359 --- /dev/null +++ b/Tests/ArrowTests/Array/BasicArrayTests.swift @@ -0,0 +1,236 @@ +// Copyright 2025 The Apache Software Foundation +// Copyright 2025 The Columnar Swift Contributors +// +// Licensed 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 Arrow +import Foundation +import Testing + +struct BasicArrayTests { + + @Test func boolArray() throws { + let builder = ArrayBuilderBoolean() + builder.append(true) + builder.appendNull() + builder.append(false) + builder.append(false) + let array = builder.finish() + #expect(array.length == 4) + #expect(array[0]! == true) + #expect(array[1] == nil) + #expect(array[2]! == false) + #expect(array[3]! == false) + #expect(array.bufferSizes == [1, 1]) + } + + @Test func uint8Array() throws { + let arrayBuilder: ArrayBuilderFixedWidth = .init() + for index: UInt8 in 0..<100 { + arrayBuilder.append(index) + } + arrayBuilder.appendNull() + #expect(arrayBuilder.length == 101) + let array = arrayBuilder.finish() + #expect(array.length == 101) + #expect(array[1] == 1) + #expect(array[10] == 10) + #expect(array[100] == nil) + for index in 0..<100 { + #expect(array[Int(index)] == UInt8(index)) + } + let slice = array.slice(offset: 5, length: 5) + for i in 0..<5 { + #expect(slice[i] == UInt8(5 + i)) + } + #expect(array.bufferSizes == [(101 + 7) / 8, 101]) + } + + @Test func int64Array() throws { + var rng = getSeededRNG() + let count = Int.random(in: 0...100_000) + var expected = [Int64](repeating: 0, count: count) + for i in 0.. = .init() + for i in 0...stride]) + } + + @Test func stringArray() throws { + let builder: ArrayBuilderVariableLength = .init() + builder.appendNull() + builder.append("abc") + builder.append("def") + builder.appendNull() + builder.append("This is a longer string") + builder.appendNull() + builder.appendNull() + for i in 0..<100 { + builder.append("test \(i)") + } + let array = builder.finish() + #expect(array[0] == nil) + #expect(array[1] == "abc") + #expect(array[2] == "def") + #expect(array[3] == nil) + #expect(array[4] == "This is a longer string") + #expect(array[5] == nil) + #expect(array[6] == nil) + var utf8Count = 3 + 3 + 23 + for i in 0..<100 { + utf8Count += "test \(i)".utf8.count + #expect(array[i + 7] == "test \(i)") + } + #expect(array.length == 107) + #expect(array.bufferSizes == [(107 + 7) / 8, (107 + 1) * 4, utf8Count]) + } + + @Test func doubleArray() throws { + let builder: ArrayBuilderFixedWidth = .init() + builder.append(14) + builder.appendNull() + builder.append(40.4) + let array = builder.finish() + #expect(array.length == 3) + #expect(array.nullCount == 1) + #expect(array[0]! == 14) + #expect(array[1] == nil) + #expect(array[2]! == 40.4) + } + + @Test func date32Array() throws { + let date32Builder: ArrayBuilderDate32 = .init() + let date2 = Date(timeIntervalSinceReferenceDate: 86400 * 1) + let date1 = Date(timeIntervalSinceReferenceDate: 86400 * 5000 + 352) + date32Builder.append(date1) + date32Builder.append(date2) + date32Builder.appendNull() + let date32Array = date32Builder.finish() + #expect(date32Array.length == 3) + #expect(date32Array[1] == date2) + let adjustedDate1 = Date( + timeIntervalSince1970: date1.timeIntervalSince1970 - 352) + #expect(date32Array[0]! == adjustedDate1) + } + + @Test func date64Array() throws { + let date64Builder: ArrayBuilderDate64 = .init() + let date2 = Date(timeIntervalSinceReferenceDate: 86400 * 1) + let date1 = Date(timeIntervalSinceReferenceDate: 86400 * 5000 + 352) + date64Builder.append(date1) + date64Builder.append(date2) + date64Builder.appendNull() + let date64Array = date64Builder.finish() + #expect(date64Array.length == 3) + #expect(date64Array[1] == date2) + #expect(date64Array[0]! == date1) + } + + @Test func time32Array() throws { + let milliBuilder: ArrayBuilderTime32 = .init() + milliBuilder.append(100) + milliBuilder.append(1_000_000) + milliBuilder.appendNull() + let milliArray = milliBuilder.finish() + #expect(milliArray.length == 3) + #expect(milliArray[1] == 1_000_000) + #expect(milliArray[2] == nil) + + let secBuilder: ArrayBuilderTime32 = .init() + secBuilder.append(200) + secBuilder.appendNull() + secBuilder.append(2_000_011) + let secArray = secBuilder.finish() + #expect(secArray.length == 3) + #expect(secArray[1] == nil) + #expect(secArray[2] == 2_000_011) + } + + @Test func time64Array() throws { + let nanoBuilder: ArrayBuilderTime64 = .init() + nanoBuilder.append(10000) + nanoBuilder.appendNull() + nanoBuilder.append(123_456_789) + let nanoArray = nanoBuilder.finish() + #expect(nanoArray.length == 3) + #expect(nanoArray[1] == nil) + #expect(nanoArray[2] == 123_456_789) + + let microBuilder: ArrayBuilderTime64 = .init() + microBuilder.appendNull() + microBuilder.append(20000) + microBuilder.append(987_654_321) + + let microArray = microBuilder.finish() + #expect(microArray.length == 3) + #expect(microArray[1] == 20000) + #expect(microArray[2] == 987_654_321) + } + + @Test func timestampArray() throws { + // Test timestamp with seconds unit + let secBuilder: ArrayBuilderTimestamp = .init() + secBuilder.append(1_609_459_200) // 2021-01-01 00:00:00 + secBuilder.append(1_609_545_600) // 2021-01-02 00:00:00 + secBuilder.appendNull() + let secArray = secBuilder.finish() + #expect(secArray.length == 3) + #expect(secArray[0] == 1_609_459_200) + #expect(secArray[1] == 1_609_545_600) + #expect(secArray[2] == nil) + + // Test timestamp with milliseconds unit and timezone America/New_York + let msBuilder: ArrayBuilderTimestamp = .init() + msBuilder.append(1_609_459_200_000) // 2021-01-01 00:00:00.000 + msBuilder.appendNull() + msBuilder.append(1_609_545_600_000) // 2021-01-02 00:00:00.000 + let msArray = msBuilder.finish() + #expect(msArray.length == 3) + #expect(msArray[0] == 1_609_459_200_000) + #expect(msArray[1] == nil) + #expect(msArray[2] == 1_609_545_600_000) + + // Test timestamp with microseconds unit and timezone UTC + let usBuilder: ArrayBuilderTimestamp = .init() + usBuilder.append(1_609_459_200_000_000) // 2021-01-01 00:00:00.000000 + usBuilder.append(1_609_545_600_000_000) // 2021-01-02 00:00:00.000000 + usBuilder.append(1_609_632_000_000_000) // 2021-01-03 00:00:00.000000 + let usArray = usBuilder.finish() + #expect(usArray.length == 3) + #expect(usArray[0] == 1_609_459_200_000_000) + #expect(usArray[1] == 1_609_545_600_000_000) + #expect(usArray[2] == 1_609_632_000_000_000) + + // Test timestamp with nanoseconds unit + let nsBuilder: ArrayBuilderTimestamp = .init() + nsBuilder.appendNull() + // 2021-01-01 00:00:00.000000000 + nsBuilder.append(1_609_459_200_000_000_000) + // 2021-01-02 00:00:00.000000000 + nsBuilder.append(1_609_545_600_000_000_000) + let nsArray = nsBuilder.finish() + #expect(nsArray.length == 3) + #expect(nsArray[0] == nil) + #expect(nsArray[1] == 1_609_459_200_000_000_000) + #expect(nsArray[2] == 1_609_545_600_000_000_000) + } +} diff --git a/Tests/ArrowTests/Array/FuzzedArrayTests.swift b/Tests/ArrowTests/Array/FuzzedArrayTests.swift new file mode 100644 index 0000000..2e9169d --- /dev/null +++ b/Tests/ArrowTests/Array/FuzzedArrayTests.swift @@ -0,0 +1,335 @@ +// Copyright 2025 The Apache Software Foundation +// Copyright 2025 The Columnar Swift Contributors +// +// Licensed 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 Arrow +import Foundation +import Testing + +struct FuzzedArrayTests { + + @Test func stringArrayWithRandomNulls() throws { + var rng = getSeededRNG() + let count = Int.random(in: 0...100_000) + var testArray = [String?](repeating: nil, count: count) + // Random strings with random nulls + var utf8Count: Int = 0 + var nullCount: Int = 0 + for i in 0.. = .init() + for value in testArray { + if let value { + builder.append(value) + } else { + builder.appendNull() + } + } + let array = builder.finish() + #expect(array.length == count) + for i in 0.. = .init() + var byteCount: Int = 0 + let count: Int = 100 + var nullCount: Int = 0 + for index in 0.. = .init() + for value in expected { + if let value { + builder.append(value) + } else { + builder.appendNull() + } + } + let array = builder.finish() + #expect(array.length == count) + #expect(array.nullCount == nullCount) + for i in 0.. = .init() + for value in expected { + if let value { + builder.append(value) + } else { + builder.appendNull() + } + } + let array = builder.finish() + #expect(array.length == count) + #expect(array.nullCount == nullCount) + for i in 0...stride, + ] + #expect(array.bufferSizes == expectedBufferSizes) + } + + @Test func stringArrayVaryingNullDensity() throws { + var rng = getSeededRNG() + let densities = [0.0, 0.1, 0.5, 0.9, 1.0] + for nullProbability in densities { + let count = Int.random(in: 0...10_000) + var byteCount: Int = 0 + var nullCount: Int = 0 + var expected = [String?](repeating: nil, count: count) + for i in 0.. nullProbability { + let length = Int.random(in: 0...50, using: &rng) + expected[i] = randomString(length: length, using: &rng) + } + } + let arrayBuilder: ArrayBuilderVariableLength = .init() + for value in expected { + if let value { + arrayBuilder.append(value) + byteCount += value.utf8.count + } else { + arrayBuilder.appendNull() + nullCount += 1 + } + } + let array = arrayBuilder.finish() + #expect(array.length == count) + #expect(array.nullCount == nullCount) + for i in 0.. = .init() + for value in expected { + if let value { + builder.append(value) + } else { + builder.appendNull() + nullCount += 1 + } + } + let array = builder.finish() + #expect(array.length == count) + #expect(array.nullCount == nullCount) + for i in 0.. = .init() + for value in expected { + if let value { + builder.append(value) + } else { + builder.appendNull() + nullCount += 1 + } + } + let array = builder.finish() + #expect(array.length == count) + #expect(array.nullCount == nullCount) + for i in 0...stride, + ] + #expect(array.bufferSizes == expectedBufferSizes) + } +} diff --git a/Tests/ArrowTests/Array/ReadmeExamples.swift b/Tests/ArrowTests/Array/ReadmeExamples.swift new file mode 100644 index 0000000..a6d1515 --- /dev/null +++ b/Tests/ArrowTests/Array/ReadmeExamples.swift @@ -0,0 +1,79 @@ +// Copyright 2025 The Columnar Swift Contributors +// +// Licensed 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 +import Testing + +@testable import Arrow + +struct ReadmeExamples { + + @Test func int8Array() throws { + let swiftArray: [Int8?] = [1, nil, 2, 3, nil, 4] + let arrayBuilder: ArrayBuilderFixedWidth = .init() + for value in swiftArray { + if let value { + arrayBuilder.append(value) + } else { + arrayBuilder.appendNull() + } + } + let arrowArray = arrayBuilder.finish() + for i in 0.. = .init() + for value in swiftArray { + if let value { + arrayBuilder.append(value) + } else { + arrayBuilder.appendNull() + } + } + let arrowArray = arrayBuilder.finish() + #expect(arrowArray[0] == "ab") + #expect(arrowArray[1] == nil) + #expect(arrowArray[2] == "c") + #expect(arrowArray[3] == "") + #expect(arrowArray[4] == ".") + + // for i in 0..() - let nameBuilder = ArrayBuilderVariable() + let nameBuilder = ArrayBuilderVariableLength() // Create struct builder let structBuilder = ArrayBuilderStruct(fields: [ diff --git a/Tests/ArrowTests/BufferTests.swift b/Tests/ArrowTests/BufferTests.swift index 5d500eb..3cd126d 100644 --- a/Tests/ArrowTests/BufferTests.swift +++ b/Tests/ArrowTests/BufferTests.swift @@ -62,7 +62,7 @@ struct BufferTests { } let buffer = builder.finish() - #expect(buffer.length == 10000) + #expect(buffer.length == 10_000 * MemoryLayout.stride) } @Test func fixedWidthBufferTinyInitialCapacity() throws { diff --git a/Tests/ArrowTests/CDataTests.swift b/Tests/ArrowTests/CDataTests.swift index 9a44011..c7389ea 100644 --- a/Tests/ArrowTests/CDataTests.swift +++ b/Tests/ArrowTests/CDataTests.swift @@ -120,7 +120,7 @@ struct CDataTests { let importer = ArrowCImporter() switch importer.importArray(UnsafePointer(cArrayMutPtr), arrowType: .utf8) { case .success(let holder): - let builder = RecordBatch.Builder() + let builder = RecordBatchX.Builder() switch builder .addColumn("test", arrowArray: holder) .finish() diff --git a/Tests/ArrowTests/CodableTests.swift b/Tests/ArrowTests/CodableTests.swift index 8ae2f0f..7397c30 100644 --- a/Tests/ArrowTests/CodableTests.swift +++ b/Tests/ArrowTests/CodableTests.swift @@ -90,7 +90,7 @@ struct CodableTests { doubleBuilder.append(101.1, nil, nil) stringBuilder.append("test0", "test1", "test2") dateBuilder.append(date1, date1, date1) - let result = RecordBatch.Builder() + let result = RecordBatchX.Builder() .addColumn("propBool", arrowArray: try boolBuilder.finish()) .addColumn("propInt8", arrowArray: try int8Builder.finish()) .addColumn("propInt16", arrowArray: try int16Builder.finish()) @@ -138,7 +138,7 @@ struct CodableTests { let int8Builder: NumberArrayBuilder = try ArrowArrayBuilders.loadNumberArrayBuilder() int8Builder.append(10, 11, 12) - let result = RecordBatch.Builder() + let result = RecordBatchX.Builder() .addColumn("propInt8", arrowArray: try int8Builder.finish()) .finish() switch result { @@ -158,7 +158,7 @@ struct CodableTests { let int8WNilBuilder: NumberArrayBuilder = try ArrowArrayBuilders.loadNumberArrayBuilder() int8WNilBuilder.append(10, nil, 12, nil) - let resultWNil = RecordBatch.Builder() + let resultWNil = RecordBatchX.Builder() .addColumn( "propInt8", arrowArray: try int8WNilBuilder.finish() @@ -187,7 +187,7 @@ struct CodableTests { let stringBuilder = try ArrowArrayBuilders.loadStringArrayBuilder() int8Builder.append(10, 11, 12, 13) stringBuilder.append("test10", "test11", "test12", "test13") - switch RecordBatch.Builder() + switch RecordBatchX.Builder() .addColumn("propInt8", arrowArray: try int8Builder.finish()) .addColumn("propString", arrowArray: try stringBuilder.finish()) .finish() @@ -202,7 +202,7 @@ struct CodableTests { throw err } - switch RecordBatch.Builder() + switch RecordBatchX.Builder() .addColumn("propString", arrowArray: try stringBuilder.finish()) .addColumn("propInt8", arrowArray: try int8Builder.finish()) .finish() @@ -224,7 +224,7 @@ struct CodableTests { let stringWNilBuilder = try ArrowArrayBuilders.loadStringArrayBuilder() int8Builder.append(10, 11, 12, 13) stringWNilBuilder.append(nil, "test11", nil, "test13") - let resultWNil = RecordBatch.Builder() + let resultWNil = RecordBatchX.Builder() .addColumn("propInt8", arrowArray: try int8Builder.finish()) .addColumn("propString", arrowArray: try stringWNilBuilder.finish()) .finish() @@ -245,7 +245,7 @@ struct CodableTests { } } - func getArrayValue(_ rb: RecordBatch, colIndex: Int, rowIndex: UInt) -> T? + func getArrayValue(_ rb: RecordBatchX, colIndex: Int, rowIndex: UInt) -> T? { let anyArray = rb.columns[colIndex] return anyArray.asAny(UInt(rowIndex)) as? T diff --git a/Tests/ArrowTests/IPCTests.swift b/Tests/ArrowTests/IPCTests.swift index 8b0b9ea..7faa41c 100644 --- a/Tests/ArrowTests/IPCTests.swift +++ b/Tests/ArrowTests/IPCTests.swift @@ -51,9 +51,9 @@ func loadArrowResource(name: String) throws(ArrowError) -> URL { @discardableResult func checkBoolRecordBatch( - _ result: Result -) throws(ArrowError) -> [RecordBatch] { - let recordBatches: [RecordBatch] + _ result: Result +) throws(ArrowError) -> [RecordBatchX] { + let recordBatches: [RecordBatchX] switch result { case .success(let result): recordBatches = result.batches @@ -89,9 +89,9 @@ func checkBoolRecordBatch( @discardableResult func checkStructRecordBatch( - _ result: Result -) throws(ArrowError) -> [RecordBatch] { - let recordBatches: [RecordBatch] + _ result: Result +) throws(ArrowError) -> [RecordBatchX] { + let recordBatches: [RecordBatchX] switch result { case .success(let result): recordBatches = result.batches @@ -165,7 +165,7 @@ func makeStructSchema() throws -> ArrowSchema { .finish() } -func makeStructRecordBatch() throws -> RecordBatch { +func makeStructRecordBatch() throws -> RecordBatchX { let testData = StructTest() let dateNow = Date.now let structBuilder = try ArrowArrayBuilders.structArrayBuilderForType( @@ -183,7 +183,7 @@ func makeStructRecordBatch() throws -> RecordBatch { Float(22.22), "23", Data("24".utf8), dateNow, ]) let structArray = try structBuilder.finish() - let result = RecordBatch.Builder() + let result = RecordBatchX.Builder() .addColumn("struct1", arrowArray: structArray) .finish() switch result { @@ -194,7 +194,7 @@ func makeStructRecordBatch() throws -> RecordBatch { } } -func makeRecordBatch() throws -> RecordBatch { +func makeRecordBatch() throws -> RecordBatchX { let uint8Builder: NumberArrayBuilder = try ArrowArrayBuilders.loadNumberArrayBuilder() uint8Builder.append(10) @@ -231,7 +231,7 @@ func makeRecordBatch() throws -> RecordBatch { let date32Array = try date32Builder.finish() let int32Array = try int32Builder.finish() let floatArray = try floatBuilder.finish() - let result = RecordBatch.Builder() + let result = RecordBatchX.Builder() .addColumn("col1", arrowArray: uint8Array) .addColumn("col2", arrowArray: stringArray) .addColumn("col3", arrowArray: date32Array) @@ -256,7 +256,7 @@ struct IPCStreamReaderTests { .recordbatch, schema: schema, batches: [recordBatch]) switch arrowWriter.writeStreaming(writerInfo) { case .success(let writeData): - let arrowReader = ArrowReader() + let arrowReader = ArrowReaderX() switch arrowReader.readStreaming(writeData) { case .success(let result): let recordBatches = result.batches @@ -300,9 +300,9 @@ struct IPCStreamReaderTests { struct IPCFileReaderTests { @Test func fileReader_double() throws { let fileURL = try loadArrowResource(name: "testdata_double") - let arrowReader = ArrowReader() + let arrowReader = ArrowReaderX() let result = arrowReader.fromFile(fileURL) - let recordBatches: [RecordBatch] + let recordBatches: [RecordBatchX] switch result { case .success(let result): recordBatches = result.batches @@ -335,14 +335,14 @@ struct IPCFileReaderTests { @Test func fileReader_bool() throws { let fileURL = try loadArrowResource(name: "testdata_bool") - let arrowReader = ArrowReader() + let arrowReader = ArrowReaderX() try checkBoolRecordBatch(arrowReader.fromFile(fileURL)) } @Test func fileWriter_bool() throws { // read existing file let fileURL = try loadArrowResource(name: "testdata_bool") - let arrowReader = ArrowReader() + let arrowReader = ArrowReaderX() let fileRBs = try checkBoolRecordBatch(arrowReader.fromFile(fileURL)) let arrowWriter = ArrowWriter() // write data from file to a stream @@ -368,14 +368,14 @@ struct IPCFileReaderTests { @Test func fileReader_struct() throws { let fileURL = try loadArrowResource(name: "testdata_struct") - let arrowReader = ArrowReader() + let arrowReader = ArrowReaderX() try checkStructRecordBatch(arrowReader.fromFile(fileURL)) } @Test func fileWriter_struct() throws { // read existing file let fileURL = try loadArrowResource(name: "testdata_struct") - let arrowReader = ArrowReader() + let arrowReader = ArrowReaderX() let fileRBs = try checkStructRecordBatch(arrowReader.fromFile(fileURL)) let arrowWriter = ArrowWriter() // write data from file to a stream @@ -408,7 +408,7 @@ struct IPCFileReaderTests { .recordbatch, schema: schema, batches: [recordBatch]) switch arrowWriter.writeFile(writerInfo) { case .success(let writeData): - let arrowReader = ArrowReader() + let arrowReader = ArrowReaderX() switch arrowReader.readFile(writeData) { case .success(let result): let recordBatches = result.batches @@ -455,7 +455,7 @@ struct IPCFileReaderTests { let writerInfo = ArrowWriter.Info(.schema, schema: schema) switch arrowWriter.writeFile(writerInfo) { case .success(let writeData): - let arrowReader = ArrowReader() + let arrowReader = ArrowReaderX() switch arrowReader.readFile(writeData) { case .success(let result): #expect(result.schema != nil) @@ -479,7 +479,7 @@ struct IPCFileReaderTests { } } - func makeBinaryDataset() throws -> (ArrowSchema, RecordBatch) { + func makeBinaryDataset() throws -> (ArrowSchema, RecordBatchX) { let schemaBuilder = ArrowSchema.Builder() let schema = schemaBuilder.addField( "binary", type: .binary, isNullable: false @@ -493,7 +493,7 @@ struct IPCFileReaderTests { binaryBuilder.append("test44".data(using: .utf8)) let binaryArray = try binaryBuilder.finish() - let result = RecordBatch.Builder() + let result = RecordBatchX.Builder() .addColumn("binary", arrowArray: binaryArray) .finish() switch result { @@ -504,7 +504,7 @@ struct IPCFileReaderTests { } } - func makeTimeDataset() throws -> (ArrowSchema, RecordBatch) { + func makeTimeDataset() throws -> (ArrowSchema, RecordBatchX) { let schemaBuilder = ArrowSchema.Builder() let schema = schemaBuilder.addField( "time64", type: .time64(.microsecond), isNullable: false @@ -528,7 +528,7 @@ struct IPCFileReaderTests { time32Builder.append(3) let time64Array = try time64Builder.finish() let time32Array = try time32Builder.finish() - let result = RecordBatch.Builder() + let result = RecordBatchX.Builder() .addColumn("time64", arrowArray: time64Array) .addColumn("time32", arrowArray: time32Array) .finish() @@ -552,7 +552,7 @@ struct IPCFileReaderTests { ) switch arrowWriter.writeStreaming(writerInfo) { case .success(let writeData): - let arrowReader = ArrowReader() + let arrowReader = ArrowReaderX() switch arrowReader.readStreaming(writeData) { case .success(let result): let recordBatches = result.batches @@ -611,7 +611,7 @@ struct IPCFileReaderTests { let arrowWriter = ArrowWriter() switch arrowWriter.writeFile(writerInfo) { case .success(let writeData): - let arrowReader = ArrowReader() + let arrowReader = ArrowReaderX() switch arrowReader.readFile(writeData) { case .success(let result): #expect(result.schema != nil) @@ -644,7 +644,7 @@ struct IPCFileReaderTests { let arrowWriter = ArrowWriter() switch arrowWriter.writeFile(writerInfo) { case .success(let writeData): - let arrowReader = ArrowReader() + let arrowReader = ArrowReaderX() switch arrowReader.readFile(writeData) { case .success(let result): #expect(result.schema != nil) diff --git a/Tests/ArrowTests/RecordBatchTests.swift b/Tests/ArrowTests/RecordBatchTests.swift index 120f633..647f104 100644 --- a/Tests/ArrowTests/RecordBatchTests.swift +++ b/Tests/ArrowTests/RecordBatchTests.swift @@ -30,7 +30,7 @@ struct RecordBatchTests { let intArray = try uint8Builder.finish() let stringArray = try stringBuilder.finish() - let result = RecordBatch.Builder() + let result = RecordBatchX.Builder() .addColumn("col1", arrowArray: intArray) .addColumn("col2", arrowArray: stringArray) .finish() @@ -64,7 +64,7 @@ struct RecordBatchTests { let array = try stringBuilder.finish() let field = ArrowField(name: "col1", dataType: .utf8, isNullable: false) - let result = RecordBatch.Builder() + let result = RecordBatchX.Builder() .addColumn(field, arrowArray: array) .finish() if case .success(_) = result { diff --git a/Tests/ArrowTests/TableTests.swift b/Tests/ArrowTests/TableTests.swift index acf5795..59aa61b 100644 --- a/Tests/ArrowTests/TableTests.swift +++ b/Tests/ArrowTests/TableTests.swift @@ -207,7 +207,7 @@ struct TableTests { stringBuilder.append("test22") let intArray = try uint8Builder.finish() let stringArray = try stringBuilder.finish() - let result = RecordBatch.Builder() + let result = RecordBatchX.Builder() .addColumn("col1", arrowArray: intArray) .addColumn("col2", arrowArray: stringArray) .finish().flatMap({ rb in