From ae63723dfe9027ab1a9eb92058168134aec5e816 Mon Sep 17 00:00:00 2001 From: Will Temperley Date: Sat, 31 Jan 2026 12:21:05 +0800 Subject: [PATCH 1/2] Dicionary array added. --- Sources/Arrow/Array/Builder.swift | 12 +- Sources/Arrow/Array/DictionaryArray.swift | 68 ++++ Tests/ArrowTests/ArrayNullBufferTests.swift | 4 +- Tests/ArrowTests/BasicArrayTests.swift | 4 +- Tests/ArrowTests/ChunkedArrayTests.swift | 4 +- Tests/ArrowTests/CodableTests.swift | 385 -------------------- Tests/ArrowTests/DictionaryArrayTests.swift | 63 ++++ Tests/ArrowTests/FuzzedArrayTests.swift | 6 +- Tests/ArrowTests/ListArrayTests.swift | 2 +- Tests/ArrowTests/ReadmeExamples.swift | 2 +- Tests/ArrowTests/RecordBatchTests.swift | 2 +- Tests/ArrowTests/StructArrayTests.swift | 2 +- Tests/ArrowTests/TableTests.swift | 8 +- 13 files changed, 154 insertions(+), 408 deletions(-) create mode 100644 Sources/Arrow/Array/DictionaryArray.swift delete mode 100644 Tests/ArrowTests/CodableTests.swift create mode 100644 Tests/ArrowTests/DictionaryArrayTests.swift diff --git a/Sources/Arrow/Array/Builder.swift b/Sources/Arrow/Array/Builder.swift index c099b88..2b35f2c 100644 --- a/Sources/Arrow/Array/Builder.swift +++ b/Sources/Arrow/Array/Builder.swift @@ -75,7 +75,7 @@ public class ArrayBuilderBoolean: AnyArrayBuilder { } /// A builder for Arrow arrays holding fixed-width types. -public class ArrayBuilderFixedWidth: +public class ArrayBuilderNumeric: AnyArrayBuilder { @@ -234,7 +234,7 @@ typealias ArrayBuilderBinary = ArrayBuilderVariableLength /// A builder for Arrow arrays holding `Date`s with a resolution of one day. public struct ArrayBuilderDate32: AnyArrayBuilder { public typealias ArrayType = ArrowArrayDate32 - let builder: ArrayBuilderFixedWidth = .init() + let builder: ArrayBuilderNumeric = .init() public init() {} @@ -259,7 +259,7 @@ public struct ArrayBuilderDate32: AnyArrayBuilder { /// A builder for Arrow arrays holding `Date`s with a resolution of one day. public struct ArrayBuilderDate64: AnyArrayBuilder { public typealias ArrayType = ArrowArrayDate64 - let builder: ArrayBuilderFixedWidth = .init() + let builder: ArrayBuilderNumeric = .init() public init() {} @@ -282,13 +282,13 @@ public struct ArrayBuilderDate64: AnyArrayBuilder { } /// A builder for Arrow arrays holding Time32 values. -public typealias ArrayBuilderTime32 = ArrayBuilderFixedWidth +public typealias ArrayBuilderTime32 = ArrayBuilderNumeric /// A builder for Arrow arrays holding Time64 values. -public typealias ArrayBuilderTime64 = ArrayBuilderFixedWidth +public typealias ArrayBuilderTime64 = ArrayBuilderNumeric /// A builder for Arrow arrays holding Timestamp values. -public typealias ArrayBuilderTimestamp = ArrayBuilderFixedWidth +public typealias ArrayBuilderTimestamp = ArrayBuilderNumeric public class ArrayBuilderList { diff --git a/Sources/Arrow/Array/DictionaryArray.swift b/Sources/Arrow/Array/DictionaryArray.swift new file mode 100644 index 0000000..c055a7b --- /dev/null +++ b/Sources/Arrow/Array/DictionaryArray.swift @@ -0,0 +1,68 @@ +// Copyright 2026 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. + +/// A shared container for the dictionary values. +public class DictionaryValues: @unchecked Sendable { + var currentArray: AnyArrowArrayProtocol + + init(_ array: AnyArrowArrayProtocol) { + self.currentArray = array + } +} + +///// An Arrow dictionary array. +public struct ArrowDictionaryArray< + IndexType: FixedWidthInteger & BitwiseCopyable +>: ArrowArrayProtocol { + public let offset: Int + public let length: Int + public var bufferSizes: [Int] { keys.bufferSizes } + public var buffers: [ArrowBufferProtocol] { keys.buffers } + public var nullCount: Int { keys.nullCount } + + public let keys: ArrowArrayNumeric + public let values: DictionaryValues + + public init( + offset: Int = 0, + length: Int, + keys: ArrowArrayNumeric, + values: AnyArrowArrayProtocol + ) { + self.offset = offset + self.length = length + self.keys = keys + self.values = DictionaryValues(values) + } + + public subscript(index: Int) -> Any? { + precondition(index >= 0 && index < length, "Invalid index.") + let offsetIndex = self.offset + index + guard let key = keys[offsetIndex] else { + return nil + } + precondition( + Int(key) < values.currentArray.length, "Key out of bounds for dictionary") + return values.currentArray.any(at: Int(key)) + } + + public func slice(offset: Int, length: Int) -> Self { + .init( + offset: 0, + length: length, + keys: keys.slice(offset: self.offset + offset, length: length), + values: values.currentArray + ) + } +} diff --git a/Tests/ArrowTests/ArrayNullBufferTests.swift b/Tests/ArrowTests/ArrayNullBufferTests.swift index 60c33dc..e0a1eac 100644 --- a/Tests/ArrowTests/ArrayNullBufferTests.swift +++ b/Tests/ArrowTests/ArrayNullBufferTests.swift @@ -21,7 +21,7 @@ struct ArrayNullBufferTests { @Test func allValidValues() throws { // Should be able to omit null buffer entirely - let arrayBuilder: ArrayBuilderFixedWidth = .init() + let arrayBuilder: ArrayBuilderNumeric = .init() for i in 0..<1000 { arrayBuilder.append(Int64(i)) // No nulls } @@ -35,7 +35,7 @@ struct ArrayNullBufferTests { } @Test func allNullValues() throws { - let arrayBuilder: ArrayBuilderFixedWidth = .init() + let arrayBuilder: ArrayBuilderNumeric = .init() for _ in 0..<1000 { arrayBuilder.appendNull() } diff --git a/Tests/ArrowTests/BasicArrayTests.swift b/Tests/ArrowTests/BasicArrayTests.swift index 66eecbc..62cb457 100644 --- a/Tests/ArrowTests/BasicArrayTests.swift +++ b/Tests/ArrowTests/BasicArrayTests.swift @@ -38,7 +38,7 @@ struct BasicArrayTests { } @Test func uint8Array() throws { - let arrayBuilder: ArrayBuilderFixedWidth = .init() + let arrayBuilder: ArrayBuilderNumeric = .init() for index: UInt8 in 0..<100 { arrayBuilder.append(index) } @@ -105,7 +105,7 @@ struct BasicArrayTests { } @Test func doubleArray() throws { - let builder: ArrayBuilderFixedWidth = .init() + let builder: ArrayBuilderNumeric = .init() builder.append(14) builder.appendNull() builder.append(40.4) diff --git a/Tests/ArrowTests/ChunkedArrayTests.swift b/Tests/ArrowTests/ChunkedArrayTests.swift index a45cabe..61e598b 100644 --- a/Tests/ArrowTests/ChunkedArrayTests.swift +++ b/Tests/ArrowTests/ChunkedArrayTests.swift @@ -38,7 +38,7 @@ struct ChunkedArrayTests { : (remaining <= numChunks - chunks.count) ? 1 : Int.random(in: 1...(remaining - (numChunks - chunks.count - 1))) - let builder = ArrayBuilderFixedWidth() + let builder = ArrayBuilderNumeric() for i in 0..] = [] for chunkIdx in 0..<10 { - let builder = ArrayBuilderFixedWidth() + let builder = ArrayBuilderNumeric() for i in 0.. = -// try ArrowArrayBuilders.loadNumberArrayBuilder() -// let int16Builder: NumberArrayBuilder = -// try ArrowArrayBuilders.loadNumberArrayBuilder() -// let int32Builder: NumberArrayBuilder = -// try ArrowArrayBuilders.loadNumberArrayBuilder() -// let int64Builder: NumberArrayBuilder = -// try ArrowArrayBuilders.loadNumberArrayBuilder() -// let uint8Builder: NumberArrayBuilder = -// try ArrowArrayBuilders.loadNumberArrayBuilder() -// let uint16Builder: NumberArrayBuilder = -// try ArrowArrayBuilders.loadNumberArrayBuilder() -// let uint32Builder: NumberArrayBuilder = -// try ArrowArrayBuilders.loadNumberArrayBuilder() -// let uint64Builder: NumberArrayBuilder = -// try ArrowArrayBuilders.loadNumberArrayBuilder() -// let floatBuilder: NumberArrayBuilder = -// try ArrowArrayBuilders.loadNumberArrayBuilder() -// let doubleBuilder: NumberArrayBuilder = -// try ArrowArrayBuilders.loadNumberArrayBuilder() -// let stringBuilder = try ArrowArrayBuilders.loadStringArrayBuilder() -// let dateBuilder = try ArrowArrayBuilders.loadDate64ArrayBuilder() -// -// boolBuilder.append(false, true, false) -// int8Builder.append(10, 11, 12) -// int16Builder.append(20, 21, 22) -// int32Builder.append(30, 31, 32) -// int64Builder.append(40, 41, 42) -// uint8Builder.append(50, 51, 52) -// uint16Builder.append(60, 61, 62) -// uint32Builder.append(70, 71, 72) -// uint64Builder.append(80, 81, 82) -// floatBuilder.append(90.1, 91.1, 92.1) -// doubleBuilder.append(101.1, nil, nil) -// stringBuilder.append("test0", "test1", "test2") -// dateBuilder.append(date1, date1, date1) -// let result = RecordBatchX.Builder() -// .addColumn("propBool", arrowArray: try boolBuilder.finish()) -// .addColumn("propInt8", arrowArray: try int8Builder.finish()) -// .addColumn("propInt16", arrowArray: try int16Builder.finish()) -// .addColumn("propInt32", arrowArray: try int32Builder.finish()) -// .addColumn("propInt64", arrowArray: try int64Builder.finish()) -// .addColumn("propUInt8", arrowArray: try uint8Builder.finish()) -// .addColumn("propUInt16", arrowArray: try uint16Builder.finish()) -// .addColumn("propUInt32", arrowArray: try uint32Builder.finish()) -// .addColumn("propUInt64", arrowArray: try uint64Builder.finish()) -// .addColumn("propFloat", arrowArray: try floatBuilder.finish()) -// .addColumn("propDouble", arrowArray: try doubleBuilder.finish()) -// .addColumn("propString", arrowArray: try stringBuilder.finish()) -// .addColumn("propDate", arrowArray: try dateBuilder.finish()) -// .finish() -// switch result { -// case .success(let rb): -// let decoder = ArrowDecoder(rb) -// let testClasses = try decoder.decode(TestClass.self) -// for index in 0.. = -// try ArrowArrayBuilders.loadNumberArrayBuilder() -// int8Builder.append(10, 11, 12) -// let result = RecordBatchX.Builder() -// .addColumn("propInt8", arrowArray: try int8Builder.finish()) -// .finish() -// switch result { -// case .success(let rb): -// let decoder = ArrowDecoder(rb) -// let testData = try decoder.decode(Int8?.self) -// for index in 0.. = -// try ArrowArrayBuilders.loadNumberArrayBuilder() -// int8WNilBuilder.append(10, nil, 12, nil) -// let resultWNil = RecordBatchX.Builder() -// .addColumn( -// "propInt8", -// arrowArray: try int8WNilBuilder.finish() -// ) -// .finish() -// switch resultWNil { -// case .success(let rb): -// let decoder = ArrowDecoder(rb) -// let testData = try decoder.decode(Int8?.self) -// for index in 0.. = -// try ArrowArrayBuilders.loadNumberArrayBuilder() -// let stringBuilder = try ArrowArrayBuilders.loadStringArrayBuilder() -// int8Builder.append(10, 11, 12, 13) -// stringBuilder.append("test10", "test11", "test12", "test13") -// switch RecordBatchX.Builder() -// .addColumn("propInt8", arrowArray: try int8Builder.finish()) -// .addColumn("propString", arrowArray: try stringBuilder.finish()) -// .finish() -// { -// case .success(let rb): -// let decoder = ArrowDecoder(rb) -// let testData = try decoder.decode([Int8: String].self) -// for data in testData { -// #expect("test\(data.key)" == data.value) -// } -// case .failure(let err): -// throw err -// } -// -// switch RecordBatchX.Builder() -// .addColumn("propString", arrowArray: try stringBuilder.finish()) -// .addColumn("propInt8", arrowArray: try int8Builder.finish()) -// .finish() -// { -// case .success(let rb): -// let decoder = ArrowDecoder(rb) -// let testData = try decoder.decode([String: Int8].self) -// for data in testData { -// #expect("test\(data.value)" == data.key) -// } -// case .failure(let err): -// throw err -// } -// } -// -// @Test func arrowMapDecoderWithNull() throws { -// let int8Builder: NumberArrayBuilder = -// try ArrowArrayBuilders.loadNumberArrayBuilder() -// let stringWNilBuilder = try ArrowArrayBuilders.loadStringArrayBuilder() -// int8Builder.append(10, 11, 12, 13) -// stringWNilBuilder.append(nil, "test11", nil, "test13") -// let resultWNil = RecordBatchX.Builder() -// .addColumn("propInt8", arrowArray: try int8Builder.finish()) -// .addColumn("propString", arrowArray: try stringWNilBuilder.finish()) -// .finish() -// switch resultWNil { -// case .success(let rb): -// let decoder = ArrowDecoder(rb) -// let testData = try decoder.decode([Int8: String?].self) -// for data in testData { -// let str = data.value -// if data.key % 2 == 0 { -// #expect(str == nil) -// } else { -// #expect(str == "test\(data.key)") -// } -// } -// case .failure(let err): -// throw err -// } -// } -// -// func getArrayValue(_ rb: RecordBatch, colIndex: Int, rowIndex: UInt) -> T? -// { -// let anyArray = rb.columns[colIndex] -// return anyArray.asAny(UInt(rowIndex)) as? T -// } -// -// @Test func arrowKeyedEncoder() throws { -// var infos: [TestClass] = [] -// for index in 0..<10 { -// let tClass = TestClass() -// let offset = index * 12 -// tClass.propBool = index % 2 == 0 -// tClass.propInt8 = Int8(offset + 1) -// tClass.propInt16 = Int16(offset + 2) -// tClass.propInt32 = Int32(offset + 3) -// tClass.propInt64 = Int64(offset + 4) -// tClass.propUInt8 = UInt8(offset + 5) -// tClass.propUInt16 = UInt16(offset + 6) -// tClass.propUInt32 = UInt32(offset + 7) -// tClass.propUInt64 = UInt64(offset + 8) -// tClass.propFloat = Float(offset + 9) -// tClass.propDouble = index % 2 == 0 ? Double(offset + 10) : nil -// tClass.propString = "\(offset + 11)" -// tClass.propDate = Date.now -// infos.append(tClass) -// } -// -// let rb = try ArrowEncoder.encode(infos)! -// #expect(Int(rb.length) == infos.count) -// #expect(rb.columns.count == 13) -// #expect(rb.columns[0].type == .boolean) -// #expect(rb.columns[1].type == .int8) -// #expect(rb.columns[2].type == .int16) -// #expect(rb.columns[3].type == .int32) -// #expect(rb.columns[4].type == .int64) -// #expect(rb.columns[5].type == .uint8) -// #expect(rb.columns[6].type == .uint16) -// #expect(rb.columns[7].type == .uint32) -// #expect(rb.columns[8].type == .uint64) -// #expect(rb.columns[9].type == .float32) -// #expect(rb.columns[10].type == .float64) -// #expect(rb.columns[11].type == .utf8) -// #expect(rb.columns[12].type == .date64) -// for index in 0..<10 { -// let offset = index * 12 -// #expect( -// getArrayValue(rb, colIndex: 0, rowIndex: UInt(index)) -// == (index % 2 == 0)) -// #expect( -// getArrayValue(rb, colIndex: 1, rowIndex: UInt(index)) -// == Int8(offset + 1)) -// #expect( -// getArrayValue(rb, colIndex: 2, rowIndex: UInt(index)) -// == Int16(offset + 2) -// ) -// #expect( -// getArrayValue(rb, colIndex: 3, rowIndex: UInt(index)) -// == Int32(offset + 3) -// ) -// #expect( -// getArrayValue(rb, colIndex: 4, rowIndex: UInt(index)) -// == Int64(offset + 4) -// ) -// #expect( -// getArrayValue(rb, colIndex: 5, rowIndex: UInt(index)) -// == UInt8(offset + 5) -// ) -// #expect( -// getArrayValue(rb, colIndex: 6, rowIndex: UInt(index)) -// == UInt16(offset + 6)) -// #expect( -// getArrayValue(rb, colIndex: 7, rowIndex: UInt(index)) -// == UInt32(offset + 7)) -// #expect( -// getArrayValue(rb, colIndex: 8, rowIndex: UInt(index)) -// == UInt64(offset + 8)) -// #expect( -// getArrayValue(rb, colIndex: 9, rowIndex: UInt(index)) -// == Float(offset + 9) -// ) -// if index % 2 == 0 { -// #expect( -// getArrayValue(rb, colIndex: 10, rowIndex: UInt(index)) -// == Double(offset + 10)) -// } else { -// #expect( -// getArrayValue(rb, colIndex: 10, rowIndex: UInt(index)) == Double?(nil) -// ) -// } -// #expect( -// getArrayValue(rb, colIndex: 11, rowIndex: UInt(index)) -// == String(offset + 11)) -// } -// } -// -// @Test func arrowUnkeyedEncoder() throws { -// var testMap: [Int8: String?] = [:] -// for index in 0..<10 { -// testMap[Int8(index)] = "test\(index)" -// } -// -// let rb = try ArrowEncoder.encode(testMap) -// #expect(Int(rb.length) == testMap.count) -// #expect(rb.columns.count == 2) -// #expect(rb.columns[0].type == .int8) -// #expect(rb.columns[1].type == .utf8) -// for index in 0..<10 { -// let key: Int8 = getArrayValue(rb, colIndex: 0, rowIndex: UInt(index))! -// let value: String = getArrayValue(rb, colIndex: 1, rowIndex: UInt(index))! -// #expect("test\(key)" == value) -// } -// } -// -// @Test func arrowSingleEncoder() throws { -// var intArray: [Int32?] = [] -// for index in 0..<100 { -// if index == 10 { -// intArray.append(nil) -// } else { -// intArray.append(Int32(index)) -// } -// } -// -// let rb = try ArrowEncoder.encode(intArray)! -// #expect(Int(rb.length) == intArray.count) -// #expect(rb.columns.count == 1) -// #expect(rb.columns[0].type == .int32) -// for index in 0..<100 { -// if index == 10 { -// let anyArray = rb.columns[0] -// #expect(anyArray.asAny(UInt(index)) == nil) -// } else { -// #expect( -// getArrayValue(rb, colIndex: 0, rowIndex: UInt(index)) == Int32(index)) -// } -// } -// } -//} diff --git a/Tests/ArrowTests/DictionaryArrayTests.swift b/Tests/ArrowTests/DictionaryArrayTests.swift new file mode 100644 index 0000000..6ac3e9e --- /dev/null +++ b/Tests/ArrowTests/DictionaryArrayTests.swift @@ -0,0 +1,63 @@ +// DictionaryArrayTests.swift +// Arrow +// +// Created by Will Temperley on 30/01/2026. All rights reserved. +// Copyright 2026 Will Temperley. +// +// Copying or reproduction of this file via any medium requires prior express +// written permission from the copyright holder. +// ----------------------------------------------------------------------------- +/// +/// Implementation notes, links and internal documentation go here. +/// +// ----------------------------------------------------------------------------- + +import Foundation +import Testing + +@testable import Arrow + +struct DictionaryArrayTests { + + @Test func basicArrayTests() throws { + + let keyArrayBuilder: ArrayBuilderNumeric = .init() + keyArrayBuilder.append(0) + keyArrayBuilder.append(1) + keyArrayBuilder.append(2) + keyArrayBuilder.append(3) + keyArrayBuilder.append(4) + let keyArray = keyArrayBuilder.finish() + + let valueArrayBuilder: ArrayBuilderVariableLength = .init() + valueArrayBuilder.append("A") + valueArrayBuilder.append("B") + valueArrayBuilder.append("C") + valueArrayBuilder.append("D") + valueArrayBuilder.append("E") + let valuesArray = valueArrayBuilder.finish() + + let dictionaryArray = ArrowDictionaryArray( + length: 5, + keys: keyArray, + values: valuesArray + ) + + #expect(dictionaryArray[2] as? String == "C") + #expect(dictionaryArray[3] as? String == "D") + + let valueArrayBuilder2: ArrayBuilderVariableLength = .init() + valueArrayBuilder2.append("F") + valueArrayBuilder2.append("G") + valueArrayBuilder2.append("H") + valueArrayBuilder2.append("I") + valueArrayBuilder2.appendNull() + let valuesArray2 = valueArrayBuilder2.finish() + dictionaryArray.values.currentArray = valuesArray2 + + #expect(dictionaryArray[2] as? String == "H") + #expect(dictionaryArray[3] as? String == "I") + #expect(dictionaryArray[4] == nil) + } + +} diff --git a/Tests/ArrowTests/FuzzedArrayTests.swift b/Tests/ArrowTests/FuzzedArrayTests.swift index ebada6b..9d53240 100644 --- a/Tests/ArrowTests/FuzzedArrayTests.swift +++ b/Tests/ArrowTests/FuzzedArrayTests.swift @@ -26,7 +26,7 @@ struct FuzzedArrayTests { for i in 0.. = .init() + let arrayBuilder: ArrayBuilderNumeric = .init() for i in 0.. = .init() + let builder: ArrayBuilderNumeric = .init() for value in expected { if let value { builder.append(value) @@ -325,7 +325,7 @@ struct FuzzedArrayTests { } i += runLength } - let builder: ArrayBuilderFixedWidth = .init() + let builder: ArrayBuilderNumeric = .init() for value in expected { if let value { builder.append(value) diff --git a/Tests/ArrowTests/ListArrayTests.swift b/Tests/ArrowTests/ListArrayTests.swift index a79e3ac..fef05e7 100644 --- a/Tests/ArrowTests/ListArrayTests.swift +++ b/Tests/ArrowTests/ListArrayTests.swift @@ -21,7 +21,7 @@ struct ListArrayTests { @Test func int32Example() { let builder = ArrayBuilderList( - valueBuilder: ArrayBuilderFixedWidth()) + valueBuilder: ArrayBuilderNumeric()) builder.append { childBuilder in childBuilder.append(1) diff --git a/Tests/ArrowTests/ReadmeExamples.swift b/Tests/ArrowTests/ReadmeExamples.swift index 41ed43e..519aa3f 100644 --- a/Tests/ArrowTests/ReadmeExamples.swift +++ b/Tests/ArrowTests/ReadmeExamples.swift @@ -21,7 +21,7 @@ struct ReadmeExamples { @Test func int8Array() throws { let swiftArray: [Int8?] = [1, nil, 2, 3, nil, 4] - let arrayBuilder: ArrayBuilderFixedWidth = .init() + let arrayBuilder: ArrayBuilderNumeric = .init() for value in swiftArray { if let value { arrayBuilder.append(value) diff --git a/Tests/ArrowTests/RecordBatchTests.swift b/Tests/ArrowTests/RecordBatchTests.swift index 937e7b4..7808825 100644 --- a/Tests/ArrowTests/RecordBatchTests.swift +++ b/Tests/ArrowTests/RecordBatchTests.swift @@ -18,7 +18,7 @@ import Testing struct RecordBatchTests { @Test func recordBatch() throws { - let uint8Builder = ArrayBuilderFixedWidth() + let uint8Builder = ArrayBuilderNumeric() uint8Builder.append(10) uint8Builder.append(22) uint8Builder.appendNull() diff --git a/Tests/ArrowTests/StructArrayTests.swift b/Tests/ArrowTests/StructArrayTests.swift index 97c4be3..d93e81a 100644 --- a/Tests/ArrowTests/StructArrayTests.swift +++ b/Tests/ArrowTests/StructArrayTests.swift @@ -21,7 +21,7 @@ struct StructArrayTests { @Test func testStructArray() { // Create builders for struct fields - let idBuilder = ArrayBuilderFixedWidth() + let idBuilder = ArrayBuilderNumeric() let nameBuilder = ArrayBuilderVariableLength() // Create struct builder diff --git a/Tests/ArrowTests/TableTests.swift b/Tests/ArrowTests/TableTests.swift index 87a3386..a40fcae 100644 --- a/Tests/ArrowTests/TableTests.swift +++ b/Tests/ArrowTests/TableTests.swift @@ -38,7 +38,7 @@ struct TableTests { } @Test func table() throws { - let doubleBuilder: ArrayBuilderFixedWidth = .init() + let doubleBuilder: ArrayBuilderNumeric = .init() doubleBuilder.append(11.11) doubleBuilder.append(22.22) let stringBuilder = ArrayBuilderVariableLength() @@ -86,12 +86,12 @@ struct TableTests { } @Test func tableWithChunkedData() throws { - let uint8Builder: ArrayBuilderFixedWidth = .init() + let uint8Builder: ArrayBuilderNumeric = .init() uint8Builder.append(10) uint8Builder.append(22) - let uint8Builder2: ArrayBuilderFixedWidth = .init() + let uint8Builder2: ArrayBuilderNumeric = .init() uint8Builder2.append(33) - let uint8Builder3: ArrayBuilderFixedWidth = .init() + let uint8Builder3: ArrayBuilderNumeric = .init() uint8Builder3.append(44) let stringBuilder = ArrayBuilderVariableLength() stringBuilder.append("test10") From 14aeed78765e7a03794b22a630311446f5708501 Mon Sep 17 00:00:00 2001 From: Will Temperley Date: Tue, 24 Feb 2026 10:26:20 +0800 Subject: [PATCH 2/2] Data sources menu fixed.. --- Package.swift | 2 +- Sources/Arrow/ArrowField.swift | 33 +++++++------- Sources/Arrow/ArrowType.swift | 13 +++--- Sources/ArrowIPC/ArrowReader.swift | 45 +++++++++++++++++++ Sources/ArrowIPC/FlatBuffersTypeAliases.swift | 10 ++++- Tests/ArrowIPCTests/ArrowTestingGold.swift | 4 +- 6 files changed, 78 insertions(+), 29 deletions(-) diff --git a/Package.swift b/Package.swift index a487c86..d76d08b 100644 --- a/Package.swift +++ b/Package.swift @@ -20,7 +20,7 @@ import PackageDescription let package = Package( name: "Arrow", platforms: [ - .macOS(.v26) + .macOS(.v26), .iOS(.v26), .watchOS(.v26), .tvOS(.v26), .visionOS(.v26), ], products: [ .library(name: "Arrow", targets: ["Arrow"]), diff --git a/Sources/Arrow/ArrowField.swift b/Sources/Arrow/ArrowField.swift index 600dea2..b5409de 100644 --- a/Sources/Arrow/ArrowField.swift +++ b/Sources/Arrow/ArrowField.swift @@ -31,7 +31,6 @@ public struct ArrowField: Codable, Sendable { /// /// If true, the field *may* contain null values. public var isNullable: Bool - public let orderedDict: Bool /// A map of key-value pairs containing additional custom meta data. public var metadata: [String: String] } @@ -60,7 +59,6 @@ extension ArrowField { self.name = name self.type = dataType self.isNullable = isNullable - self.orderedDict = false self.metadata = metadata } @@ -76,21 +74,20 @@ extension ArrowField { ) } - /// Create a new `ArrowField` suitable for `ArrowType::Dictionary`. - /// - public init( - dictWithName: String, - key: ArrowType, - value: ArrowType, - isNullable: Bool - ) { - precondition( - key.isDictionaryKeyType, - "\(key) is not a valid dictionary key" - ) - let dataType: ArrowType = .dictionary(key, value) - self = Self(name: dictWithName, dataType: dataType, isNullable: isNullable) - } +// /// Create a new `ArrowField` suitable for `ArrowType::Dictionary`. +// public init( +// dictWithName: String, +// key: ArrowType, +// value: ArrowType, +// isNullable: Bool +// ) { +// precondition( +// key.isDictionaryKeyType, +// "\(key) is not a valid dictionary key" +// ) +// let dataType: ArrowType = .dictionary(key, value) +// self = Self(name: dictWithName, dataType: dataType, isNullable: isNullable) +// } /// Create a new struct `ArrowField`. /// @@ -218,7 +215,7 @@ extension ArrowField { @inlinable public var dictIsOrdered: Bool { switch self.type { - case .dictionary: return self.orderedDict + case .dictionary(_, let isOrdered, _, _): return isOrdered default: return false } } diff --git a/Sources/Arrow/ArrowType.swift b/Sources/Arrow/ArrowType.swift index b2f9078..1e741f8 100644 --- a/Sources/Arrow/ArrowType.swift +++ b/Sources/Arrow/ArrowType.swift @@ -293,7 +293,7 @@ public indirect enum ArrowType: Codable, Sendable, Equatable { /// /// This type mostly used to represent low cardinality string /// arrays or a limited set of primitive types as integers. - case dictionary(ArrowType, ArrowType) + case dictionary(id: Int64, isOrdered: Bool, key: ArrowType, value: ArrowType) /// Exact 32-bit width decimal value with precision and scale /// /// * precision is the total number of digits @@ -496,8 +496,8 @@ extension ArrowType: CustomStringConvertible { return "LargeListView(\(elementType))" case .union(let mode, let fields): return "Union(\(mode), \(fields) fields)" - case .dictionary(let keyType, let valueType): - return "Dictionary(\(keyType), \(valueType))" + case .dictionary(let id, let isOrdered, let keyType, let valueType): + return "Dictionary(\(id), \(isOrdered), \(keyType), \(valueType))" case .decimal32(let precision, let scale): return "Decimal32(\(precision), \(scale))" case .decimal64(let precision, let scale): @@ -627,7 +627,7 @@ extension ArrowType { @inlinable public var isNested: Bool { switch self { - case .dictionary(_, let v): + case .dictionary(_, _, _, let v): return v.isNested case .runEndEncoded(_, let v): return v.type.isNested @@ -699,7 +699,8 @@ extension ArrowType { && aField.type.equalsDataType(bField.type) && aSorted == bSorted // Dictionary - case (.dictionary(let aKey, let aValue), .dictionary(let bKey, let bValue)): + case (.dictionary(_, _, let aKey, let aValue), .dictionary(_, _, let bKey, let bValue)): + // Ignoring dictionary id here. return aKey.equalsDataType(bKey) && aValue.equalsDataType(bValue) // RunEndEncoded @@ -855,7 +856,7 @@ extension ArrowType { } // Dictionary - case (.dictionary(let k1, let v1), .dictionary(let k2, let v2)): + case (.dictionary(_, _, let k1, let v1), .dictionary(_, _, let k2, let v2)): return k1.contains(k2) && v1.contains(v2) // Base case: equality diff --git a/Sources/ArrowIPC/ArrowReader.swift b/Sources/ArrowIPC/ArrowReader.swift index da27672..e273582 100644 --- a/Sources/ArrowIPC/ArrowReader.swift +++ b/Sources/ArrowIPC/ArrowReader.swift @@ -90,6 +90,51 @@ public struct ArrowReader { let arrowSchema = try Self.loadSchema(schema: schema) var recordBatches: [RecordBatch] = [] + for block in footer.dictionaries { + let (message, bodyOffset) = try data.withParserSpan { input in + try input.seek(toAbsoluteOffset: block.offset) + let marker = try UInt32(parsingLittleEndian: &input) + if marker != continuationMarker { + throw ArrowError(.invalid("Missing continuation marker.")) + } + let messageLength = try UInt32(parsingLittleEndian: &input) + let data = try [UInt8](parsing: &input, byteCount: Int(messageLength)) + var mbb = ByteBuffer(data: Data(data)) + let message: FMessage = getRoot(byteBuffer: &mbb) + let offset = Int64(input.startPosition) + return (message, offset) + } + + guard message.headerType == .dictionarybatch else { + throw ArrowError(.invalid("Expected DictionaryBatch message.")) + } + + guard let dictMessage = message.header(type: FDictionaryBatch.self) else { + throw ArrowError(.invalid("Expected DictionaryBatch as message header")) + } + + // 1. Get the Dictionary ID and 'isDelta' flag + let dictId = dictMessage.id + let isDelta = dictMessage.isDelta + + // 2. The dictionary data is actually just a RecordBatch with ONE column + // The schema for this internal batch is defined by the dictionary type + // found in the global Schema for this specific ID. + guard let rbMessage = dictMessage.data else { + throw ArrowError(.invalid("DictionaryBatch has no data")) + } + + let dictBatch = try Self.loadRecordBatch( + data: self.data, + arrowSchema: arrowSchema, + rbMessage: rbMessage, + offset: bodyOffset + ) + + // 4. Update the "Box" in your provider + // try dictionaryProvider.update(id: dictId, array: dictionaryArray, isDelta: isDelta) + } + // MARK: Record batch parsing for block in footer.recordBatches { diff --git a/Sources/ArrowIPC/FlatBuffersTypeAliases.swift b/Sources/ArrowIPC/FlatBuffersTypeAliases.swift index 9a450d4..6497e6f 100644 --- a/Sources/ArrowIPC/FlatBuffersTypeAliases.swift +++ b/Sources/ArrowIPC/FlatBuffersTypeAliases.swift @@ -13,16 +13,22 @@ // limitations under the License. typealias FFooter = org_apache_arrow_flatbuf_Footer + +typealias FMessageHeader = org_apache_arrow_flatbuf_MessageHeader typealias FMessage = org_apache_arrow_flatbuf_Message + typealias FBlock = org_apache_arrow_flatbuf_Block typealias FField = org_apache_arrow_flatbuf_Field typealias FSchema = org_apache_arrow_flatbuf_Schema typealias FBuffer = org_apache_arrow_flatbuf_Buffer typealias FFieldNode = org_apache_arrow_flatbuf_FieldNode -typealias FRecordBatch = org_apache_arrow_flatbuf_RecordBatch -typealias FMessageHeader = org_apache_arrow_flatbuf_MessageHeader typealias FKeyValue = org_apache_arrow_flatbuf_KeyValue +// MARK: Record batches. +typealias FRecordBatch = org_apache_arrow_flatbuf_RecordBatch +typealias FDictionaryBatch = org_apache_arrow_flatbuf_DictionaryBatch +typealias FDictionaryEncoding = org_apache_arrow_flatbuf_DictionaryEncoding + // MARK: Top level type. typealias FType = org_apache_arrow_flatbuf_Type_ diff --git a/Tests/ArrowIPCTests/ArrowTestingGold.swift b/Tests/ArrowIPCTests/ArrowTestingGold.swift index 108f6d4..74a9567 100644 --- a/Tests/ArrowIPCTests/ArrowTestingGold.swift +++ b/Tests/ArrowIPCTests/ArrowTestingGold.swift @@ -82,10 +82,10 @@ struct ArrowTestingGold { "generated_union", ] - @Test +// @Test func readFile() throws { - let name = "generated_nested_large_offsets" + let name = "generated_dictionary" let (testFile, testCase) = try loadTestCase( name: name, fileExtension: "arrow_file") let arrowReader = try ArrowReader(url: testFile)