Merge pull request #23 from willtemperley/main

Add struct array support to ArrowIPC.

Author: willtemperley

Sources/Arrow/Array/Array.swift

@@ -14,7 +14,7 @@
 
 import Foundation
 
-public protocol ArrowArrayProtocol<ItemType> {
+public protocol ArrowArrayProtocol {
   associatedtype ItemType
   subscript(_ index: Int) -> ItemType? { get }
   var offset: Int { get }
@@ -70,18 +70,32 @@ public struct ArrowArrayBoolean: ArrowArrayProtocol {
 }
 
 /// An Arrow array of fixed-width types.
-struct ArrowArrayFixed<Element, ValueBuffer>: ArrowArrayProtocol
+public struct ArrowArrayFixed<ValueBuffer>: ArrowArrayProtocol
 where
-  Element: Numeric,
-  ValueBuffer: FixedWidthBufferProtocol<Element>
+  ValueBuffer: FixedWidthBufferProtocol,
+  ValueBuffer.ElementType: Numeric
 {
-  typealias ItemType = Element
-  let offset: Int
-  let length: Int
+  public typealias ItemType = ValueBuffer.ElementType
+
+  //  public typealias ItemType = Element
+  public let offset: Int
+  public let length: Int
   let nullBuffer: NullBuffer
   let valueBuffer: ValueBuffer
 
-  subscript(index: Int) -> Element? {
+  public init(
+    offset: Int = 0,
+    length: Int,
+    nullBuffer: NullBuffer,
+    valueBuffer: ValueBuffer
+  ) {
+    self.offset = offset
+    self.length = length
+    self.nullBuffer = nullBuffer
+    self.valueBuffer = valueBuffer
+  }
+
+  public subscript(index: Int) -> ValueBuffer.ElementType? {
     precondition(index >= 0 && index < length, "Invalid index.")
     let offsetIndex = self.offset + index
     if !self.nullBuffer.isSet(offsetIndex) {
@@ -90,7 +104,7 @@ where
     return valueBuffer[offsetIndex]
   }
 
-  func slice(offset: Int, length: Int) -> Self {
+  public func slice(offset: Int, length: Int) -> Self {
     .init(
       offset: offset,
       length: length,
@@ -116,7 +130,7 @@ where
   let valueBuffer: ValueBuffer
 
   public init(
-    offset: Int,
+    offset: Int = 0,
     length: Int,
     nullBuffer: NullBuffer,
     offsetsBuffer: OffsetsBuffer,
@@ -162,7 +176,7 @@ where
 {
   typealias ItemType = Date
 
-  let array: ArrowArrayFixed<Date32, ValueBuffer>
+  let array: ArrowArrayFixed<ValueBuffer>
 
   var offset: Int {
     array.offset
@@ -192,11 +206,11 @@ where
 /// An Arrow array of `Date`s with a resolution of 1 second.
 struct ArrowArrayDate64<ValueBuffer>: ArrowArrayProtocol
 where
-  ValueBuffer: FixedWidthBufferProtocol<Int64>
+  ValueBuffer: FixedWidthBufferProtocol<Date64>
 {
   typealias ItemType = Date
 
-  let array: ArrowArrayFixed<Date64, ValueBuffer>
+  let array: ArrowArrayFixed<ValueBuffer>
 
   var offset: Int {
     array.offset
@@ -259,3 +273,44 @@ where
     )
   }
 }
+
+/// 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 init(
+    offset: Int = 0,
+    length: Int,
+    nullBuffer: NullBuffer,
+    fields: [(name: String, array: any ArrowArrayProtocol)]
+  ) {
+    self.offset = offset
+    self.length = length
+    self.nullBuffer = nullBuffer
+    self.fields = fields
+  }
+
+  public subscript(index: Int) -> ItemType? {
+    guard nullBuffer.isSet(offset + index) else { return nil }
+
+    var result: [String: Any] = [:]
+    for (name, array) in fields {
+      result[name] = array.any(at: index)
+    }
+    return result
+  }
+
+  public func slice(offset newOffset: Int, length newLength: Int) -> Self {
+    .init(
+      offset: self.offset + newOffset,
+      length: newLength,
+      nullBuffer: nullBuffer,
+      fields: fields
+    )
+  }
+}

Sources/Arrow/Array/Builder.swift

@@ -76,7 +76,7 @@ class ArrayBuilderBoolean: AnyArrayBuilder {
 /// A builder for Arrow arrays holding fixed-width types.
 class ArrayBuilderFixedWidth<T: Numeric>: AnyArrayBuilder {
 
-  typealias ArrayType = ArrowArrayFixed<T, FixedWidthBuffer<T>>
+  typealias ArrayType = ArrowArrayFixed<FixedWidthBuffer<T>>
 
   var length: Int
   let nullBuilder: NullBufferBuilder

Sources/Arrow/Array/StructArray.swift

@@ -14,8 +14,6 @@
 
 /// A type which provides offset ranges in Arrow arrays.
 protocol OffsetsBuffer {
-  /// Number of offset pairs available
-  //  var count: Int { get }
 
   /// Get the start and end offsets for the element at index
   /// - Parameter index: Zero-based index of the element

Sources/Arrow/Buffer/OffsetProtocol.swift

@@ -15,6 +15,7 @@
 import Arrow
 import Foundation
 
+/// A `Data` backed Arrow utf8 array.
 typealias ArrowArrayUtf8 = ArrowArrayVariable<
   String,
   FixedWidthBufferIPC<Int32>,
@@ -39,7 +40,6 @@ extension ArrowArrayUtf8 {
     let offsetsBufferTyped = FixedWidthBufferIPC<Int32>(buffer: offsetsBuffer)
     let valueBufferTyped = VariableLengthBufferIPC<String>(buffer: valueBuffer)
     return Self(
-      offset: 0,
       length: length,
       nullBuffer: nullBuffer,
       offsetsBuffer: offsetsBufferTyped,
@@ -72,7 +72,6 @@ extension ArrowArrayBinary {
     let offsetsBufferTyped = FixedWidthBufferIPC<Int32>(buffer: offsetsBuffer)
     let valueBufferTyped = VariableLengthBufferIPC<Data>(buffer: valueBuffer)
     return Self(
-      offset: 0,
       length: length,
       nullBuffer: nullBuffer,
       offsetsBuffer: offsetsBufferTyped,