apacheGH-41945: [Swift] Add interface ArrowArrayHolderBuilder

abandy · Jun 3, 2024 · f8f3278 · f8f3278
1 parent da0eb7e
commit f8f3278
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 2 deletions.
diff --git a/swift/Arrow/Sources/Arrow/ArrowArrayBuilder.swift b/swift/Arrow/Sources/Arrow/ArrowArrayBuilder.swift
@@ -17,7 +17,12 @@
 
 import Foundation
 
-public class ArrowArrayBuilder<T: ArrowBufferBuilder, U: ArrowArray<T.ItemType>> {
+public protocol ArrowArrayHolderBuilder {
+    func toHolder() throws -> ArrowArrayHolder
+    func appendAny(_ val: Any?)
+}
+
+public class ArrowArrayBuilder<T: ArrowBufferBuilder, U: ArrowArray<T.ItemType>>: ArrowArrayHolderBuilder {
     let type: ArrowType
     let bufferBuilder: T
     public var length: UInt {return self.bufferBuilder.length}
@@ -34,6 +39,10 @@ public class ArrowArrayBuilder<T: ArrowBufferBuilder, U: ArrowArray<T.ItemType>>
         self.bufferBuilder.append(val)
     }
 
+    public func appendAny(_ val: Any?) {
+        self.bufferBuilder.append(val as? T.ItemType)
+    }
+
     public func finish() throws -> ArrowArray<T.ItemType> {
         let buffers = self.bufferBuilder.finish()
         let arrowData = try ArrowData(self.type, buffers: buffers, nullCount: self.nullCount)
@@ -43,6 +52,10 @@ public class ArrowArrayBuilder<T: ArrowBufferBuilder, U: ArrowArray<T.ItemType>>
     public func getStride() -> Int {
         return self.type.getStride()
     }
+
+    public func toHolder() throws -> ArrowArrayHolder {
+        return try ArrowArrayHolderImpl(self.finish())
+    }
 }
 
 public class NumberArrayBuilder<T>: ArrowArrayBuilder<FixedBufferBuilder<T>, FixedArray<T>> {

diff --git a/swift/Arrow/Tests/ArrowTests/ArrayTests.swift b/swift/Arrow/Tests/ArrowTests/ArrayTests.swift
@@ -245,4 +245,31 @@ final class ArrayTests: XCTestCase {
         try checkHolderForType(ArrowType(ArrowType.ArrowBool))
         try checkHolderForType(ArrowType(ArrowType.ArrowString))
     }
-}
+
+    func testArrowArrayHolderBuilder() throws {
+        let uint8HBuilder: ArrowArrayHolderBuilder =
+            (try ArrowArrayBuilders.loadNumberArrayBuilder() as NumberArrayBuilder<UInt8>)
+        for index in 0..<100 {
+            uint8HBuilder.appendAny(UInt8(index))
+        }
+
+        let uint8Holder = try uint8HBuilder.toHolder()
+        XCTAssertEqual(uint8Holder.nullCount, 0)
+        XCTAssertEqual(uint8Holder.length, 101)
+        let uint8Data = uint8Holder.data
+
+        let stringHBuilder: ArrowArrayHolderBuilder =
+            (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.toHolder()
+        XCTAssertEqual(stringHolder.nullCount, 0)
+        XCTAssertEqual(stringHolder.length, 101)
+    }
+ }