Merge pull request #16 from willtemperley/main

Add offsets buffer.

Author: willtemperley

Sources/Arrow/ArrowArray.swift

@@ -34,7 +34,6 @@ public protocol AnyArrowArray {
 /// The interface for Arrow array types.
 public protocol ArrowArray<ItemType>: AnyArrowArray {
   associatedtype ItemType
-  var arrowData: ArrowData { get }
   init(_ arrowData: ArrowData) throws(ArrowError)
   subscript(_ index: UInt) -> ItemType? { get }
 }
@@ -132,25 +131,24 @@ public class FixedArray<T>: ArrowArrayBase<T> where T: BitwiseCopyable {
 public class StringArray: ArrowArrayBase<String> {
 
   public override subscript(_ index: UInt) -> String? {
-    let offsetIndex = MemoryLayout<Int32>.stride * Int(index)
     if self.arrowData.isNull(index) {
       return nil
     }
 
     let offsets = self.arrowData.buffers[1]
-    let values = self.arrowData.buffers[2]
-
+    let offsetIndex = MemoryLayout<Int32>.stride * Int(index)
     var startIndex: Int32 = 0
     if index > 0 {
       startIndex = offsets.rawPointer.advanced(by: offsetIndex)
         .load(as: Int32.self)
     }
-
     let endIndex = offsets.rawPointer.advanced(
       by: offsetIndex + MemoryLayout<Int32>.stride
     )
     .load(as: Int32.self)
+
     let arrayLength = Int(endIndex - startIndex)
+    let values = self.arrowData.buffers[2]
     let rawPointer = values.rawPointer.advanced(by: Int(startIndex))
       .bindMemory(to: UInt8.self, capacity: arrayLength)
     let buffer = UnsafeBufferPointer<UInt8>(

Sources/Arrow/ArrowData.swift

@@ -22,6 +22,8 @@ public struct ArrowData {
   public let length: UInt
   public let stride: Int
 
+  let nullBuffer: ArrowBuffer
+
   init(
     _ arrowType: ArrowType,
     buffers: [ArrowBuffer],
@@ -31,7 +33,8 @@ public struct ArrowData {
       arrowType, buffers: buffers,
       children: [ArrowData](),
       nullCount: nullCount,
-      length: buffers[1].length)
+      length: buffers[1].length
+    )
   }
 
   init(
@@ -47,10 +50,11 @@ public struct ArrowData {
     self.nullCount = nullCount
     self.length = length
     self.stride = arrowType.getStride()
+
+    self.nullBuffer = buffers[0]
   }
 
   public func isNull(_ at: UInt) -> Bool {
-    let nullBuffer = buffers[0]
-    return nullBuffer.length > 0 && !BitUtility.isSet(at, buffer: nullBuffer)
+    nullBuffer.length > 0 && !BitUtility.isSet(at, buffer: nullBuffer)
   }
 }

Sources/Arrow/ArrowReader.swift

@@ -20,33 +20,49 @@ let fileMarker = Data("ARROW1".utf8)
 let continuationMarker = UInt32(0xFFFF_FFFF)
 
 public struct ArrowReader: Sendable {
+
+  func makeBuffer(
+    _ buffer: FBuffer,
+    fileData: Data,
+    length: UInt,
+    messageOffset: Int64
+  ) -> ArrowBuffer {
+    let startOffset = messageOffset + buffer.offset
+    let endOffset = startOffset + buffer.length
+
+    // TODO: This should not copy.
+
+    let bufferData = [UInt8](fileData[startOffset..<endOffset])
+    return ArrowBuffer.createBuffer(bufferData, length: length)
+  }
+
   private class RecordBatchData {
-    let schema: Schema
-    let recordBatch: FlatRecordBatch
+    let schema: FSchema
+    let recordBatch: FRecordBatch
     private var fieldIndex: Int32 = 0
     private var nodeIndex: Int32 = 0
     private var bufferIndex: Int32 = 0
     init(
-      _ recordBatch: FlatRecordBatch,
-      schema: Schema
+      _ recordBatch: FRecordBatch,
+      schema: FSchema
     ) {
       self.recordBatch = recordBatch
       self.schema = schema
     }
 
-    func nextNode() -> FieldNode? {
+    func nextNode() -> FFieldNode? {
       if nodeIndex >= self.recordBatch.nodesCount { return nil }
       defer { nodeIndex += 1 }
       return self.recordBatch.nodes(at: nodeIndex)
     }
 
-    func nextBuffer() -> Buffer? {
+    func nextBuffer() -> FBuffer? {
       if bufferIndex >= self.recordBatch.buffersCount { return nil }
       defer { bufferIndex += 1 }
       return self.recordBatch.buffers(at: bufferIndex)
     }
 
-    func nextField() -> FlatField? {
+    func nextField() -> FField? {
       if fieldIndex >= self.schema.fieldsCount { return nil }
       defer { fieldIndex += 1 }
       return self.schema.fields(at: fieldIndex)
@@ -64,15 +80,15 @@ public struct ArrowReader: Sendable {
   }
 
   public class ArrowReaderResult {
-    fileprivate var messageSchema: Schema?
+    fileprivate var messageSchema: FSchema?
     public var schema: ArrowSchema?
     public var batches: [RecordBatch] = []
   }
 
   public init() {}
 
   private func loadSchema(
-    _ schema: Schema
+    _ schema: FSchema
   ) -> Result<ArrowSchema, ArrowError> {
     let builder = ArrowSchema.Builder()
     for index in 0..<schema.fieldsCount {
@@ -100,7 +116,7 @@ public struct ArrowReader: Sendable {
 
   private func loadStructData(
     _ loadInfo: DataLoadInfo,
-    field: FlatField
+    field: FField
   ) -> Result<AnyArrowArray, ArrowError> {
     guard let node = loadInfo.batchData.nextNode() else {
       return .failure(.invalid("Node not found"))
@@ -135,7 +151,7 @@ public struct ArrowReader: Sendable {
     )
   }
 
-  private func loadListData(_ loadInfo: DataLoadInfo, field: FlatField)
+  private func loadListData(_ loadInfo: DataLoadInfo, field: FField)
     -> Result<AnyArrowArray, ArrowError>
   {
     guard let node = loadInfo.batchData.nextNode() else {
@@ -166,7 +182,8 @@ public struct ArrowReader: Sendable {
     switch loadField(loadInfo, field: childField) {
     case .success(let childHolder):
       return makeArrayHolder(
-        field, buffers: [arrowNullBuffer, arrowOffsetBuffer],
+        field,
+        buffers: [arrowNullBuffer, arrowOffsetBuffer],
         nullCount: UInt(node.nullCount),
         children: [childHolder.arrowData],
         rbLength: UInt(loadInfo.batchData.recordBatch.length))
@@ -177,7 +194,7 @@ public struct ArrowReader: Sendable {
 
   private func loadPrimitiveData(
     _ loadInfo: DataLoadInfo,
-    field: FlatField
+    field: FField
   )
     -> Result<AnyArrowArray, ArrowError>
   {
@@ -206,9 +223,11 @@ public struct ArrowReader: Sendable {
       rbLength: UInt(loadInfo.batchData.recordBatch.length))
   }
 
+  // MARK: Variable data loading
+
   private func loadVariableData(
     _ loadInfo: DataLoadInfo,
-    field: FlatField
+    field: FField
   )
     -> Result<AnyArrowArray, ArrowError>
   {
@@ -232,6 +251,7 @@ public struct ArrowReader: Sendable {
     let arrowNullBuffer = makeBuffer(
       nullBuffer, fileData: loadInfo.fileData,
       length: nullLength, messageOffset: loadInfo.messageOffset)
+
     let arrowOffsetBuffer = makeBuffer(
       offsetBuffer, fileData: loadInfo.fileData,
       length: UInt(node.length), messageOffset: loadInfo.messageOffset)
@@ -246,7 +266,7 @@ public struct ArrowReader: Sendable {
 
   private func loadField(
     _ loadInfo: DataLoadInfo,
-    field: FlatField
+    field: FField
   )
     -> Result<AnyArrowArray, ArrowError>
   {
@@ -265,8 +285,8 @@ public struct ArrowReader: Sendable {
   }
 
   private func loadRecordBatch(
-    _ recordBatch: FlatRecordBatch,
-    schema: Schema,
+    _ recordBatch: FRecordBatch,
+    schema: FSchema,
     arrowSchema: ArrowSchema,
     data: Data,
     messageEndOffset: Int64
@@ -310,7 +330,7 @@ public struct ArrowReader: Sendable {
     var length = getUInt32(input, offset: offset)
     var streamData = input
     // TODO: The following assumes message order will populate schemaMessage first
-    var schemaMessage: Schema?
+    var schemaMessage: FSchema?
     while length != 0 {
       if length == continuationMarker {
         offset += Int(MemoryLayout<UInt32>.size)
@@ -325,10 +345,10 @@ public struct ArrowReader: Sendable {
         data: streamData,
         allowReadingUnalignedBuffers: true
       )
-      let message: Message = getRoot(byteBuffer: &dataBuffer)
+      let message: FMessage = getRoot(byteBuffer: &dataBuffer)
       switch message.headerType {
       case .recordbatch:
-        guard let rbMessage = message.header(type: FlatRecordBatch.self) else {
+        guard let rbMessage = message.header(type: FRecordBatch.self) else {
           return .failure(.invalid("Failed to parse RecordBatch message"))
         }
         guard let schemaMessage else {
@@ -353,7 +373,7 @@ public struct ArrowReader: Sendable {
         offset += Int(message.bodyLength + Int64(length))
         length = getUInt32(input, offset: offset)
       case .schema:
-        schemaMessage = message.header(type: Schema.self)
+        schemaMessage = message.header(type: FSchema.self)
         guard let schemaMessage else {
           return .failure(.invalid("Schema message not found"))
         }
@@ -397,7 +417,7 @@ public struct ArrowReader: Sendable {
     var footerBuffer = ByteBuffer(
       data: footerData,
       allowReadingUnalignedBuffers: useUnalignedBuffers)
-    let footer: Footer = getRoot(byteBuffer: &footerBuffer)
+    let footer: FFooter = getRoot(byteBuffer: &footerBuffer)
     guard let footerSchema = footer.schema else {
       return .failure(.invalid("Missing schema in footer"))
     }
@@ -410,7 +430,7 @@ public struct ArrowReader: Sendable {
     }
 
     for index in 0..<footer.recordBatchesCount {
-      guard let recordBatch: Block = footer.recordBatches(at: index) else {
+      guard let recordBatch: FBlock = footer.recordBatches(at: index) else {
         return .failure(.invalid("Missing record batch at index \(index)"))
       }
       var messageLength = fileData.withUnsafeBytes { rawBuffer in
@@ -436,10 +456,10 @@ public struct ArrowReader: Sendable {
       var mbb = ByteBuffer(
         data: recordBatchData,
         allowReadingUnalignedBuffers: useUnalignedBuffers)
-      let message: Message = getRoot(byteBuffer: &mbb)
+      let message: FMessage = getRoot(byteBuffer: &mbb)
       switch message.headerType {
       case .recordbatch:
-        guard let rbMessage = message.header(type: FlatRecordBatch.self) else {
+        guard let rbMessage = message.header(type: FRecordBatch.self) else {
           return .failure(.invalid("Expected RecordBatch as message header"))
         }
         guard let footerSchema = footer.schema else {
@@ -475,7 +495,9 @@ public struct ArrowReader: Sendable {
     _ fileURL: URL
   ) -> Result<ArrowReaderResult, ArrowError> {
     do {
-      let fileData = try Data(contentsOf: fileURL)
+      // TODO: implement alignment checks.
+      let fileData = try Data(contentsOf: fileURL, options: .mappedIfSafe)
+
       if !validateFileData(fileData) {
         return .failure(.ioError("Not a valid arrow file."))
       }
@@ -499,10 +521,10 @@ public struct ArrowReader: Sendable {
     var mbb = ByteBuffer(
       data: dataHeader,
       allowReadingUnalignedBuffers: useUnalignedBuffers)
-    let message: Message = getRoot(byteBuffer: &mbb)
+    let message: FMessage = getRoot(byteBuffer: &mbb)
     switch message.headerType {
     case .schema:
-      guard let sMessage = message.header(type: Schema.self) else {
+      guard let sMessage = message.header(type: FSchema.self) else {
         return .failure(.unknownError("Expected a schema but found none"))
       }
       switch loadSchema(sMessage) {
@@ -514,7 +536,7 @@ public struct ArrowReader: Sendable {
         return .failure(error)
       }
     case .recordbatch:
-      guard let rbMessage = message.header(type: FlatRecordBatch.self) else {
+      guard let rbMessage = message.header(type: FRecordBatch.self) else {
         return .failure(.invalid("Expected a RecordBatch but found none"))
       }
       // TODO: the result used here is also the return type. Ideally is would be constructed once as a struct (same issue as above)

Sources/Arrow/ArrowReaderHelper.swift

@@ -108,7 +108,7 @@ func makeNestedHolder(
 }
 
 func makeArrayHolder(
-  _ field: FlatField,
+  _ field: FField,
   buffers: [ArrowBuffer],
   nullCount: UInt,
   children: [ArrowData]?,