-
Notifications
You must be signed in to change notification settings - Fork 6
feat: Add support for List data types #39
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change | ||||||||||||||||||||||||||||||||||||||||
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
@@ -131,7 +131,7 @@ public class StructArrayBuilder: ArrowArrayBuilder<StructBufferBuilder, StructAr | |||||||||||||||||||||||||||||||||||||||||
public init(_ fields: [ArrowField], builders: [any ArrowArrayHolderBuilder]) throws { | ||||||||||||||||||||||||||||||||||||||||||
self.fields = fields | ||||||||||||||||||||||||||||||||||||||||||
self.builders = builders | ||||||||||||||||||||||||||||||||||||||||||
try super.init(ArrowNestedType(ArrowType.ArrowStruct, fields: fields)) | ||||||||||||||||||||||||||||||||||||||||||
try super.init(ArrowTypeStruct(ArrowType.ArrowStruct, fields: fields)) | ||||||||||||||||||||||||||||||||||||||||||
self.bufferBuilder.initializeTypeInfo(fields) | ||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||
|
@@ -143,7 +143,7 @@ public class StructArrayBuilder: ArrowArrayBuilder<StructBufferBuilder, StructAr | |||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||
self.builders = builders | ||||||||||||||||||||||||||||||||||||||||||
try super.init(ArrowNestedType(ArrowType.ArrowStruct, fields: fields)) | ||||||||||||||||||||||||||||||||||||||||||
try super.init(ArrowTypeStruct(ArrowType.ArrowStruct, fields: fields)) | ||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||
public override func append(_ values: [Any?]?) { | ||||||||||||||||||||||||||||||||||||||||||
|
@@ -174,6 +174,31 @@ public class StructArrayBuilder: ArrowArrayBuilder<StructBufferBuilder, StructAr | |||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||
public class ListArrayBuilder: ArrowArrayBuilder<ListBufferBuilder, ListArray> { | ||||||||||||||||||||||||||||||||||||||||||
let valueBuilder: any ArrowArrayHolderBuilder | ||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||
public override init(_ elementType: ArrowType) throws { | ||||||||||||||||||||||||||||||||||||||||||
self.valueBuilder = try ArrowArrayBuilders.loadBuilder(arrowType: elementType) | ||||||||||||||||||||||||||||||||||||||||||
try super.init(ArrowTypeList(elementType)) | ||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||
public override func append(_ values: [Any?]?) { | ||||||||||||||||||||||||||||||||||||||||||
self.bufferBuilder.append(values) | ||||||||||||||||||||||||||||||||||||||||||
if let vals = values { | ||||||||||||||||||||||||||||||||||||||||||
for val in vals { | ||||||||||||||||||||||||||||||||||||||||||
self.valueBuilder.appendAny(val) | ||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||
public override func finish() throws -> ListArray { | ||||||||||||||||||||||||||||||||||||||||||
let buffers = self.bufferBuilder.finish() | ||||||||||||||||||||||||||||||||||||||||||
let childData = try valueBuilder.toHolder().array.arrowData | ||||||||||||||||||||||||||||||||||||||||||
let arrowData = try ArrowData(self.type, buffers: buffers, children: [childData], nullCount: self.nullCount, length: self.length) | ||||||||||||||||||||||||||||||||||||||||||
return try ListArray(arrowData) | ||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||
public class ArrowArrayBuilders { | ||||||||||||||||||||||||||||||||||||||||||
public static func loadBuilder( // swiftlint:disable:this cyclomatic_complexity | ||||||||||||||||||||||||||||||||||||||||||
_ builderType: Any.Type) throws -> ArrowArrayHolderBuilder { | ||||||||||||||||||||||||||||||||||||||||||
|
@@ -290,6 +315,16 @@ public class ArrowArrayBuilders { | |||||||||||||||||||||||||||||||||||||||||
throw ArrowError.invalid("Expected arrow type for \(arrowType.id) not found") | ||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||
return try TimestampArrayBuilder(timestampType.unit) | ||||||||||||||||||||||||||||||||||||||||||
case .list: | ||||||||||||||||||||||||||||||||||||||||||
guard let listType = arrowType as? ArrowTypeList else { | ||||||||||||||||||||||||||||||||||||||||||
throw ArrowError.invalid("Expected ArrowTypeList for \(arrowType.id)") | ||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||
return try ListArrayBuilder(listType.elementType) | ||||||||||||||||||||||||||||||||||||||||||
case .strct: | ||||||||||||||||||||||||||||||||||||||||||
guard let structType = arrowType as? ArrowTypeStruct else { | ||||||||||||||||||||||||||||||||||||||||||
throw ArrowError.invalid("Expected ArrowStructType for \(arrowType.id)") | ||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||
return try StructArrayBuilder(structType.fields) | ||||||||||||||||||||||||||||||||||||||||||
Comment on lines
+318
to
+327
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Could you use
Suggested change
|
||||||||||||||||||||||||||||||||||||||||||
default: | ||||||||||||||||||||||||||||||||||||||||||
throw ArrowError.unknownType("Builder not found for arrow type: \(arrowType.id)") | ||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||
|
@@ -353,4 +388,12 @@ public class ArrowArrayBuilders { | |||||||||||||||||||||||||||||||||||||||||
public static func loadTimestampArrayBuilder(_ unit: ArrowTimestampUnit, timezone: String? = nil) throws -> TimestampArrayBuilder { | ||||||||||||||||||||||||||||||||||||||||||
return try TimestampArrayBuilder(unit, timezone: timezone) | ||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||
public static func loadStructArrayBuilder(_ fields: [ArrowField]) throws -> StructArrayBuilder { | ||||||||||||||||||||||||||||||||||||||||||
return try StructArrayBuilder(fields) | ||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||
public static func loadListArrayBuilder(_ elementType: ArrowType) throws -> ListArrayBuilder { | ||||||||||||||||||||||||||||||||||||||||||
return try ListArrayBuilder(elementType) | ||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||
} |
Original file line number | Diff line number | Diff line change | ||||||||
---|---|---|---|---|---|---|---|---|---|---|
|
@@ -338,14 +338,14 @@ public class Date64BufferBuilder: AbstractWrapperBufferBuilder<Date, Int64> { | |||||||||
|
||||||||||
public final class StructBufferBuilder: BaseBufferBuilder, ArrowBufferBuilder { | ||||||||||
public typealias ItemType = [Any?] | ||||||||||
var info: ArrowNestedType? | ||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Nested type should be able to handle a list without making an individual type for list and struct. A list should be able to be expressed as nested type with only a single field. Is there a reason you are choosing to move away from Nested types? |
||||||||||
var info: ArrowTypeStruct? | ||||||||||
public init() throws { | ||||||||||
let nulls = ArrowBuffer.createBuffer(0, size: UInt(MemoryLayout<UInt8>.stride)) | ||||||||||
super.init(nulls) | ||||||||||
} | ||||||||||
|
||||||||||
public func initializeTypeInfo(_ fields: [ArrowField]) { | ||||||||||
info = ArrowNestedType(ArrowType.ArrowStruct, fields: fields) | ||||||||||
info = ArrowTypeStruct(ArrowType.ArrowStruct, fields: fields) | ||||||||||
} | ||||||||||
|
||||||||||
public func append(_ newValue: [Any?]?) { | ||||||||||
|
@@ -379,3 +379,62 @@ public final class StructBufferBuilder: BaseBufferBuilder, ArrowBufferBuilder { | |||||||||
return [nulls] | ||||||||||
} | ||||||||||
} | ||||||||||
|
||||||||||
public class ListBufferBuilder: BaseBufferBuilder, ArrowBufferBuilder { | ||||||||||
public typealias ItemType = [Any?] | ||||||||||
var offsets: ArrowBuffer | ||||||||||
|
||||||||||
public required init() throws { | ||||||||||
self.offsets = ArrowBuffer.createBuffer(1, size: UInt(MemoryLayout<Int32>.stride)) | ||||||||||
let nulls = ArrowBuffer.createBuffer(0, size: UInt(MemoryLayout<UInt8>.stride)) | ||||||||||
super.init(nulls) | ||||||||||
self.offsets.rawPointer.storeBytes(of: Int32(0), as: Int32.self) | ||||||||||
} | ||||||||||
|
||||||||||
public func append(_ newValue: [Any?]?) { | ||||||||||
let index = UInt(self.length) | ||||||||||
self.length += 1 | ||||||||||
|
||||||||||
if length >= self.offsets.length { | ||||||||||
self.resize(length + 1) | ||||||||||
Comment on lines
+398
to
+399
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Should we use
Suggested change
|
||||||||||
} | ||||||||||
|
||||||||||
let offsetIndex = Int(index) * MemoryLayout<Int32>.stride | ||||||||||
let currentOffset = self.offsets.rawPointer.advanced(by: offsetIndex).load(as: Int32.self) | ||||||||||
|
||||||||||
if let vals = newValue { | ||||||||||
BitUtility.setBit(index + self.offset, buffer: self.nulls) | ||||||||||
let newOffset = currentOffset + Int32(vals.count) | ||||||||||
self.offsets.rawPointer.advanced(by: offsetIndex + MemoryLayout<Int32>.stride).storeBytes(of: newOffset, as: Int32.self) | ||||||||||
} else { | ||||||||||
self.nullCount += 1 | ||||||||||
BitUtility.clearBit(index + self.offset, buffer: self.nulls) | ||||||||||
self.offsets.rawPointer.advanced(by: offsetIndex + MemoryLayout<Int32>.stride).storeBytes(of: currentOffset, as: Int32.self) | ||||||||||
} | ||||||||||
} | ||||||||||
|
||||||||||
public override func isNull(_ index: UInt) -> Bool { | ||||||||||
return !BitUtility.isSet(index + self.offset, buffer: self.nulls) | ||||||||||
} | ||||||||||
|
||||||||||
public func resize(_ length: UInt) { | ||||||||||
if length > self.offsets.length { | ||||||||||
let resizeLength = resizeLength(self.offsets) | ||||||||||
var offsets = ArrowBuffer.createBuffer(resizeLength, size: UInt(MemoryLayout<Int32>.size)) | ||||||||||
var nulls = ArrowBuffer.createBuffer(resizeLength/8 + 1, size: UInt(MemoryLayout<UInt8>.size)) | ||||||||||
ArrowBuffer.copyCurrent(self.offsets, to: &offsets, len: self.offsets.capacity) | ||||||||||
ArrowBuffer.copyCurrent(self.nulls, to: &nulls, len: self.nulls.capacity) | ||||||||||
self.offsets = offsets | ||||||||||
self.nulls = nulls | ||||||||||
} | ||||||||||
} | ||||||||||
|
||||||||||
public func finish() -> [ArrowBuffer] { | ||||||||||
let length = self.length | ||||||||||
var nulls = ArrowBuffer.createBuffer(length/8 + 1, size: UInt(MemoryLayout<UInt8>.size)) | ||||||||||
var offsets = ArrowBuffer.createBuffer(length + 1, size: UInt(MemoryLayout<Int32>.size)) | ||||||||||
ArrowBuffer.copyCurrent(self.nulls, to: &nulls, len: nulls.capacity) | ||||||||||
ArrowBuffer.copyCurrent(self.offsets, to: &offsets, len: offsets.capacity) | ||||||||||
return [nulls, offsets] | ||||||||||
} | ||||||||||
} |
Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can we use a generic here: like
ListArray<T>: ArrowArray<[T?]>
?