diff --git a/Sources/Arrow/Array/Array.swift b/Sources/Arrow/Array/Array.swift index 9effb33..7c630d7 100644 --- a/Sources/Arrow/Array/Array.swift +++ b/Sources/Arrow/Array/Array.swift @@ -70,14 +70,18 @@ public struct ArrowArrayBoolean: ArrowArrayProtocol { } /// An Arrow array of fixed-width types. -struct ArrowArrayFixed: ArrowArrayProtocol where T: Numeric { - typealias ItemType = T +struct ArrowArrayFixed: ArrowArrayProtocol +where + Element: Numeric, + ValueBuffer: FixedWidthBufferProtocol +{ + typealias ItemType = Element let offset: Int let length: Int let nullBuffer: NullBuffer - let valueBuffer: FixedWidthBuffer + let valueBuffer: ValueBuffer - subscript(index: Int) -> T? { + subscript(index: Int) -> Element? { precondition(index >= 0 && index < length, "Invalid index.") let offsetIndex = self.offset + index if !self.nullBuffer.isSet(offsetIndex) { @@ -86,7 +90,7 @@ struct ArrowArrayFixed: ArrowArrayProtocol where T: Numeric { return valueBuffer[offsetIndex] } - func slice(offset: Int, length: Int) -> ArrowArrayFixed { + func slice(offset: Int, length: Int) -> Self { .init( offset: offset, length: length, @@ -97,22 +101,26 @@ struct ArrowArrayFixed: ArrowArrayProtocol where T: Numeric { } /// An Arrow array of variable-length types. -public struct ArrowArrayVariable: ArrowArrayProtocol -where T: VariableLength { - public typealias ItemType = T +public struct ArrowArrayVariable: + ArrowArrayProtocol +where + Element: VariableLength, + OffsetsBuffer: FixedWidthBufferProtocol, + ValueBuffer: VariableLengthBufferProtocol +{ + public typealias ItemType = Element public let offset: Int public let length: Int let nullBuffer: NullBuffer - let offsetsBuffer: any FixedWidthBufferProtocol - let valueBuffer: any VariableLengthBufferProtocol + let offsetsBuffer: OffsetsBuffer + let valueBuffer: ValueBuffer public init( offset: Int, length: Int, nullBuffer: NullBuffer, - offsetsBuffer: - any FixedWidthBufferProtocol, - valueBuffer: any VariableLengthBufferProtocol + offsetsBuffer: OffsetsBuffer, + valueBuffer: ValueBuffer ) { self.offset = offset self.length = length @@ -121,7 +129,7 @@ where T: VariableLength { self.valueBuffer = valueBuffer } - public subscript(index: Int) -> T? { + public subscript(index: Int) -> Element? { let offsetIndex = self.offset + index @@ -147,14 +155,14 @@ where T: VariableLength { } } -public typealias ArrowArrayUtf8 = ArrowArrayVariable -public typealias ArrowArrayBinary = ArrowArrayVariable - /// An Arrow array of `Date`s with a resolution of 1 day. -struct ArrowArrayDate32: ArrowArrayProtocol { +struct ArrowArrayDate32: ArrowArrayProtocol +where + ValueBuffer: FixedWidthBufferProtocol +{ typealias ItemType = Date - let array: ArrowArrayFixed + let array: ArrowArrayFixed var offset: Int { array.offset @@ -182,10 +190,13 @@ struct ArrowArrayDate32: ArrowArrayProtocol { } /// An Arrow array of `Date`s with a resolution of 1 second. -struct ArrowArrayDate64: ArrowArrayProtocol { +struct ArrowArrayDate64: ArrowArrayProtocol +where + ValueBuffer: FixedWidthBufferProtocol +{ typealias ItemType = Date - let array: ArrowArrayFixed + let array: ArrowArrayFixed var offset: Int { array.offset @@ -212,17 +223,20 @@ struct ArrowArrayDate64: ArrowArrayProtocol { } } -struct ArrowListArray: ArrowArrayProtocol where T: ArrowArrayProtocol { - - typealias ItemType = T +/// An Arrow list array which may be nested arbitrarily. +struct ArrowListArray: ArrowArrayProtocol +where + Element: ArrowArrayProtocol +{ + typealias ItemType = Element let offset: Int let length: Int let nullBuffer: NullBuffer let offsetsBuffer: FixedWidthBuffer - let values: T + let values: Element - subscript(index: Int) -> T? { + subscript(index: Int) -> Element? { precondition(index >= 0 && index < length, "Invalid index.") let offsetIndex = self.offset + index if !self.nullBuffer.isSet(offsetIndex) { diff --git a/Sources/Arrow/Array/Builder.swift b/Sources/Arrow/Array/Builder.swift index e72f32f..c3500df 100644 --- a/Sources/Arrow/Array/Builder.swift +++ b/Sources/Arrow/Array/Builder.swift @@ -76,7 +76,7 @@ class ArrayBuilderBoolean: AnyArrayBuilder { /// A builder for Arrow arrays holding fixed-width types. class ArrayBuilderFixedWidth: AnyArrayBuilder { - typealias ArrayType = ArrowArrayFixed + typealias ArrayType = ArrowArrayFixed> var length: Int let nullBuilder: NullBufferBuilder @@ -114,23 +114,25 @@ class ArrayBuilderFixedWidth: AnyArrayBuilder { } /// A builder for Arrow arrays holding variable length types. -class ArrayBuilderVariable: AnyArrayBuilder { - typealias ArrayType = ArrowArrayVariable +class ArrayBuilderVariable: AnyArrayBuilder { + typealias ArrayType = ArrowArrayVariable< + Element, FixedWidthBuffer, VariableLengthTypeBuffer + > var length: Int let nullBuilder: NullBufferBuilder let offsetsBuilder: FixedWidthBufferBuilder - let valueBuilder: VariableLengthTypeBufferBuilder + let valueBuilder: VariableLengthTypeBufferBuilder init() { self.length = 0 self.nullBuilder = NullBufferBuilder() self.offsetsBuilder = FixedWidthBufferBuilder() - self.valueBuilder = VariableLengthTypeBufferBuilder() + self.valueBuilder = VariableLengthTypeBufferBuilder() self.offsetsBuilder.append(Int32.zero) } - public func append(_ value: T) { + public func append(_ value: Element) { length += 1 nullBuilder.appendValid(true) let data = value.data @@ -176,7 +178,7 @@ typealias ArrayBuilderBinary = ArrayBuilderVariable /// A builder for Arrow arrays holding `Date`s with a resolution of one day. struct ArrayBuilderDate32: AnyArrayBuilder { - typealias ArrayType = ArrowArrayDate32 + typealias ArrayType = ArrowArrayDate32> let builder: ArrayBuilderFixedWidth = .init() @@ -200,7 +202,7 @@ struct ArrayBuilderDate32: AnyArrayBuilder { /// A builder for Arrow arrays holding `Date`s with a resolution of one day. struct ArrayBuilderDate64: AnyArrayBuilder { - typealias ArrayType = ArrowArrayDate64 + typealias ArrayType = ArrowArrayDate64> let builder: ArrayBuilderFixedWidth = .init() diff --git a/Sources/Arrow/Buffer/NullBuffer.swift b/Sources/Arrow/Buffer/NullBuffer.swift index 4ca0162..73b9fe1 100644 --- a/Sources/Arrow/Buffer/NullBuffer.swift +++ b/Sources/Arrow/Buffer/NullBuffer.swift @@ -48,10 +48,10 @@ public struct AllNullBuffer: NullBuffer { /// A bit-packed buffer used to represent nulls and booleans in Arrow arrays. final class BitPackedNullBuffer: NullBuffer { - var length: Int - var capacity: Int + let length: Int + let capacity: Int let ownsMemory: Bool - var buffer: UnsafePointer + let buffer: UnsafePointer init( length: Int, diff --git a/Sources/ArrowIPC/Array+IPC.swift b/Sources/ArrowIPC/Array+IPC.swift new file mode 100644 index 0000000..3d0efdd --- /dev/null +++ b/Sources/ArrowIPC/Array+IPC.swift @@ -0,0 +1,82 @@ +// 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 + +typealias ArrowArrayUtf8 = ArrowArrayVariable< + String, + FixedWidthBufferIPC, + VariableLengthBufferIPC +> + +extension ArrowArrayUtf8 { + + /// Build a `Data` backed Arrow utf8 array. + /// - Parameters: + /// - length: The array length. + /// - nullBuffer: The null buffer. + /// - offsetsBuffer: A view over file-backed data. + /// - valueBuffer: A view over file-backed data. + /// - Returns: A file-backed Arrow utf8 array. + static func utf8( + length: Int, + nullBuffer: NullBuffer, + offsetsBuffer: FileDataBuffer, + valueBuffer: FileDataBuffer + ) -> Self { + let offsetsBufferTyped = FixedWidthBufferIPC(buffer: offsetsBuffer) + let valueBufferTyped = VariableLengthBufferIPC(buffer: valueBuffer) + return Self( + offset: 0, + length: length, + nullBuffer: nullBuffer, + offsetsBuffer: offsetsBufferTyped, + valueBuffer: valueBufferTyped + ) + } +} + +typealias ArrowArrayBinary = ArrowArrayVariable< + Data, + FixedWidthBufferIPC, + VariableLengthBufferIPC +> + +extension ArrowArrayBinary { + + /// Build a `Data` backed Arrow binary array. + /// - Parameters: + /// - length: The array length. + /// - nullBuffer: The null buffer. + /// - offsetsBuffer: A view over file-backed data. + /// - valueBuffer: A view over file-backed data. + /// - Returns: A file-backed Arrow utf8 array. + static func binary( + length: Int, + nullBuffer: NullBuffer, + offsetsBuffer: FileDataBuffer, + valueBuffer: FileDataBuffer + ) -> Self { + let offsetsBufferTyped = FixedWidthBufferIPC(buffer: offsetsBuffer) + let valueBufferTyped = VariableLengthBufferIPC(buffer: valueBuffer) + return Self( + offset: 0, + length: length, + nullBuffer: nullBuffer, + offsetsBuffer: offsetsBufferTyped, + valueBuffer: valueBufferTyped + ) + } +} diff --git a/Sources/ArrowIPC/ArrowReader.swift b/Sources/ArrowIPC/ArrowReader.swift index 23343ed..08f32bb 100644 --- a/Sources/ArrowIPC/ArrowReader.swift +++ b/Sources/ArrowIPC/ArrowReader.swift @@ -21,11 +21,13 @@ import Foundation let fileMarker: [UInt8] = .init(Data("ARROW1".utf8)) let continuationMarker = UInt32(0xFFFF_FFFF) +/// A view over `Data` which backs an Arrow buffer. struct FileDataBuffer { let data: Data let range: Range } +/// A `Data` backed buffer for null bitmaps and boolean arrays. struct NullBufferIPC: NullBuffer { let buffer: FileDataBuffer @@ -36,38 +38,40 @@ struct NullBufferIPC: NullBuffer { func isSet(_ bit: Int) -> Bool { let byteIndex = bit / 8 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 } } -struct FixedWidthBufferIPC: FixedWidthBufferProtocol -where T: Numeric, T: BitwiseCopyable { +/// A `Data` backed buffer for fixed-width types. +struct FixedWidthBufferIPC: FixedWidthBufferProtocol +where Element: Numeric, Element: BitwiseCopyable { - typealias ElementType = T + typealias ElementType = Element let buffer: FileDataBuffer var length: Int { buffer.range.count } - subscript(index: Int) -> T { + subscript(index: Int) -> Element { let offsetIndex = buffer.range.lowerBound + index return buffer.data.withUnsafeBytes { rawBuffer in let sub = rawBuffer[buffer.range] - let span = Span(_unsafeBytes: sub) + let span = Span(_unsafeBytes: sub) return span[index] } } } -struct VariableLengthBuffer: VariableLengthBufferProtocol { - typealias ElementType = T +/// A `Data` backed buffer for variable-length types. +struct VariableLengthBufferIPC: + VariableLengthBufferProtocol +{ + typealias ElementType = Element let buffer: FileDataBuffer @@ -75,26 +79,22 @@ struct VariableLengthBuffer: VariableLengthBufferProtocol { buffer.range.count } - // For IPC buffer backed by Data func loadVariable( at startIndex: Int, arrayLength: Int - ) -> T { + ) -> Element { precondition(startIndex + arrayLength <= self.length) - return buffer.data.withUnsafeBytes { rawBuffer in let offsetStart = buffer.range.lowerBound + startIndex let offsetEnd = offsetStart + arrayLength let slice = rawBuffer[offsetStart.. from the slice let uint8Buffer = slice.bindMemory(to: UInt8.self) - return T(uint8Buffer) + return Element(uint8Buffer) } } } -/// This is for reading the Arrow file format. +/// A reader for the Arrow file format. /// /// The Arrow file format supports random access. The Arrow file format contains a header and footer /// around the Arrow streaming format. @@ -230,17 +230,32 @@ public struct ArrowReader { offset: 0, length: length, nullBuffer: nullBuffer, valueBuffer: valueBuffer) arrays.append(array) - } else if arrowType == .utf8 { + } else if arrowType.isVariable { let buffer1 = try nextBuffer( message: rbMessage, index: &bufferIndex, offset: offset, data: data) - let offsetsBuffer = FixedWidthBufferIPC(buffer: buffer1) let buffer2 = try nextBuffer( message: rbMessage, index: &bufferIndex, offset: offset, data: data) - let valueBuffer = VariableLengthBuffer(buffer: buffer2) - let stringArray = ArrowArrayUtf8( - offset: 0, length: length, nullBuffer: nullBuffer, - offsetsBuffer: offsetsBuffer, valueBuffer: valueBuffer) - arrays.append(stringArray) + + if arrowType == .utf8 { + let stringArray = ArrowArrayVariable.utf8( + length: length, + nullBuffer: nullBuffer, + offsetsBuffer: buffer1, + valueBuffer: buffer2 + ) + arrays.append(stringArray) + } else if arrowType == .binary { + let stringArray = ArrowArrayVariable.binary( + length: length, + nullBuffer: nullBuffer, + offsetsBuffer: buffer1, + valueBuffer: buffer2 + ) + arrays.append(stringArray) + } else { + throw ArrowError.notImplemented + } + } else { throw ArrowError.notImplemented } diff --git a/Tests/ArrowIPCTests/ArrowReaderTests.swift b/Tests/ArrowIPCTests/ArrowReaderTests.swift index 6a7b744..1342aa5 100644 --- a/Tests/ArrowIPCTests/ArrowReaderTests.swift +++ b/Tests/ArrowIPCTests/ArrowReaderTests.swift @@ -47,7 +47,9 @@ struct ArrowReaderTests { #expect(booleanColumn[3] == false) #expect(booleanColumn[4] == true) - guard let utf8Column = recordBatch.columns[1] as? ArrowArrayUtf8 else { + guard + let utf8Column = recordBatch.columns[1] as? ArrowArrayUtf8 + else { Issue.record("Failed to cast column to ArrowUtf8Array") return }