You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by do...@apache.org on 2026/09/12 04:09:58 UTC
(spark-connect-swift) branch main updated: [SPARK-59450] Support Codable in `createDataFrame` and `DataFrame.collect`
This is an automated email from the ASF dual-hosted git repository.
dongjoon-hyun pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/spark-connect-swift.git
The following commit(s) were added to refs/heads/main by this push:
new 4cb592c [SPARK-59450] Support Codable in `createDataFrame` and `DataFrame.collect`
4cb592c is described below
commit 4cb592c79c029f718ce59d7a8fa2d028e08215e2
Author: Dongjoon Hyun <do...@apache.org>
AuthorDate: Fri Sep 11 21:09:51 2026 -0700
[SPARK-59450] Support Codable in `createDataFrame` and `DataFrame.collect`
### What changes were proposed in this pull request?
This PR aims to support `Codable`-based DataFrame creation and collection in `SparkSession` and `DataFrame`, leveraging the in-tree `ArrowEncoder` and `ArrowDecoder`.
Specifically:
1. **`SparkSession.createDataFrame` with `Encodable`**:
- `createDataFrame<T: Encodable>(_ data: [T]) async throws -> DataFrame`: Automatically encodes an array of `Encodable` instances into an Apache Arrow `RecordBatch` via `ArrowEncoder`, infers the Spark DDL schema from the batch schema, and builds a `LocalRelation` or `CachedLocalRelation` (for payloads >= 1MiB).
- `createDataFrame<T: Encodable>(_ data: [T], _ schema: String)` and `createDataFrame<T: Encodable>(_ data: [T], _ schema: StructType)`: Supports explicit schema overrides and empty datasets with schema.
2. **`DataFrame.collect(as:)` with `Decodable`**:
- `collect<T: Decodable>(as type: T.Type = T.self) async throws -> [T]`: Executes the plan and directly decodes each Arrow `RecordBatch` into `[T]` using `ArrowDecoder`, eliminating intermediate untyped `Row` allocations and dictionary boxing.
3. **`ArrowEncoder` and `ArrowDecoder` enhancements**:
- Added support for Swift native `Int` (mapped to `Int64`) and `UInt` (mapped to `UInt64`) across keyed, unkeyed, and single-value containers.
- Added support for decoding `Date` and `TimestampNanos` from Arrow `Timestamp` columns (seconds, milliseconds, microseconds, and nanoseconds).
4. **Schema mapping utilities**:
- Added `DataType.init(_ arrowType: ArrowType) throws` and `StructType.init(_ arrowSchema: ArrowSchema) throws` to translate Arrow schemas to Spark SQL data types.
```swift
struct Person: Codable, Sendable, Equatable {
let name: String
let age: Int
}
let people = [Person(name: "Alice", age: 20), Person(name: "Bob", age: 25)]
// 1. Create DataFrame directly from Swift models (auto-inferred schema)
let df = try await spark.createDataFrame(people)
// 2. Collect query results directly into Swift models
let results: [Person] = try await df.filter("age >= 21").collect(as: Person.self)
```
### Why are the changes needed?
Currently, `SparkSession.createDataFrame` only accepts untyped rows `[[Sendable?]]` with an explicit DDL string or `StructType` schema, and `DataFrame.collect()` only returns untyped `[Row]`.
These new APIs provide a type-safe, Dataset-like programming experience in Swift, allowing developers to work seamlessly with native Swift `Codable` structs.
### Does this PR introduce _any_ user-facing change?
No behavior change because this PR adds only new public generic methods.
### How was this patch tested?
Pass the CIs with the newly added test cases.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Gemini 3.8 Flash (High)
Closes #552 from dongjoon-hyun/SPARK-59450.
Authored-by: Dongjoon Hyun <do...@apache.org>
Signed-off-by: Dongjoon Hyun <do...@apache.org>
---
Sources/SparkConnect/ArrowDecoder.swift | 157 +++++++++++++++++++--
Sources/SparkConnect/ArrowEncoder.swift | 51 +++++--
Sources/SparkConnect/DataFrame+Actions.swift | 25 ++++
Sources/SparkConnect/DataType.swift | 39 +++++
Sources/SparkConnect/SparkSession.swift | 92 ++++++++++--
Sources/SparkConnect/StructType.swift | 10 ++
Tests/SparkConnectTests/CreateDataFrameTests.swift | 76 ++++++++++
7 files changed, 415 insertions(+), 35 deletions(-)
diff --git a/Sources/SparkConnect/ArrowDecoder.swift b/Sources/SparkConnect/ArrowDecoder.swift
index ba9a6b7..672f4bb 100644
--- a/Sources/SparkConnect/ArrowDecoder.swift
+++ b/Sources/SparkConnect/ArrowDecoder.swift
@@ -132,6 +132,87 @@ public class ArrowDecoder: Decoder {
let array: AnyArray = try self.getCol(col)
return array.asAny(self.rbIndex) == nil
}
+
+ func decodeInt(from array: AnyArray, keyDescription: String) throws -> Int {
+ let raw = array.asAny(self.rbIndex)
+ if let val = raw as? Int64 {
+ return Int(val)
+ } else if let val = raw as? Int32 {
+ return Int(val)
+ } else if let val = raw as? Int16 {
+ return Int(val)
+ } else if let val = raw as? Int8 {
+ return Int(val)
+ } else if let val = raw as? UInt64 {
+ return Int(val)
+ } else if let val = raw as? UInt32 {
+ return Int(val)
+ } else if let val = raw as? UInt16 {
+ return Int(val)
+ } else if let val = raw as? UInt8 {
+ return Int(val)
+ }
+ throw ArrowError.invalid("Cannot decode Int for \(keyDescription)")
+ }
+
+ func decodeUInt(from array: AnyArray, keyDescription: String) throws -> UInt {
+ let raw = array.asAny(self.rbIndex)
+ if let val = raw as? UInt64 {
+ return UInt(val)
+ } else if let val = raw as? UInt32 {
+ return UInt(val)
+ } else if let val = raw as? UInt16 {
+ return UInt(val)
+ } else if let val = raw as? UInt8 {
+ return UInt(val)
+ } else if let val = raw as? Int64 {
+ return UInt(val)
+ } else if let val = raw as? Int32 {
+ return UInt(val)
+ } else if let val = raw as? Int16 {
+ return UInt(val)
+ } else if let val = raw as? Int8 {
+ return UInt(val)
+ }
+ throw ArrowError.invalid("Cannot decode UInt for \(keyDescription)")
+ }
+
+ func decodeDate(from array: AnyArray, keyDescription: String) throws -> Date {
+ if let date = array.asAny(self.rbIndex) as? Date {
+ return date
+ }
+ if let timestamp = array.asAny(self.rbIndex) as? Int64,
+ let timestampType = array.arrowData.type as? ArrowTypeTimestamp
+ {
+ switch timestampType.unit {
+ case .seconds:
+ return Date(timeIntervalSince1970: TimeInterval(timestamp))
+ case .milliseconds:
+ return Date(timeIntervalSince1970: TimeInterval(timestamp) / 1_000)
+ case .microseconds:
+ return Date(timeIntervalSince1970: TimeInterval(timestamp) / 1_000_000)
+ case .nanoseconds:
+ return Date(timeIntervalSince1970: TimeInterval(timestamp) / 1_000_000_000)
+ }
+ }
+ throw ArrowError.invalid("Cannot decode Date for \(keyDescription)")
+ }
+
+ func decodeTimestampNanos(from array: AnyArray, keyDescription: String) throws -> TimestampNanos {
+ if let timestamp = array.asAny(self.rbIndex) as? Int64,
+ let timestampType = array.arrowData.type as? ArrowTypeTimestamp
+ {
+ let epochNanos: Int64
+ switch timestampType.unit {
+ case .seconds: epochNanos = timestamp * 1_000_000_000
+ case .milliseconds: epochNanos = timestamp * 1_000_000
+ case .microseconds: epochNanos = timestamp * 1_000
+ case .nanoseconds: epochNanos = timestamp
+ }
+ return TimestampNanos(epochNanos: epochNanos)
+ }
+ throw ArrowError.invalid("Cannot decode TimestampNanos for \(keyDescription)")
+ }
}
private struct ArrowUnkeyedDecoding: UnkeyedDecodingContainer {
@@ -158,7 +239,24 @@ private struct ArrowUnkeyedDecoding: UnkeyedDecodingContainer {
}
mutating func decode<T>(_ type: T.Type) throws -> T where T: Decodable {
- if type == Int8?.self || type == Int16?.self || type == Int32?.self || type == Int64?.self
+ if type == Int.self {
+ defer { increment() }
+ let col = try self.decoder.getCol(self.currentIndex)
+ return try self.decoder.decodeInt(from: col, keyDescription: "column \(self.currentIndex)") as! T
+ } else if type == UInt.self {
+ defer { increment() }
+ let col = try self.decoder.getCol(self.currentIndex)
+ return try self.decoder.decodeUInt(from: col, keyDescription: "column \(self.currentIndex)") as! T
+ } else if type == Date.self {
+ defer { increment() }
+ let col = try self.decoder.getCol(self.currentIndex)
+ return try self.decoder.decodeDate(from: col, keyDescription: "column \(self.currentIndex)") as! T
+ } else if type == TimestampNanos.self {
+ defer { increment() }
+ let col = try self.decoder.getCol(self.currentIndex)
+ return try self.decoder.decodeTimestampNanos(
+ from: col, keyDescription: "column \(self.currentIndex)") as! T
+ } else if type == Int8?.self || type == Int16?.self || type == Int32?.self || type == Int64?.self
|| type == UInt8?.self || type == UInt16?.self || type == UInt32?.self || type == UInt64?.self
|| type == String?.self || type == Double?.self || type == Float?.self || type == Date?.self
|| type == Bool?.self || type == Bool.self || type == Int8.self || type == Int16.self
@@ -223,8 +321,8 @@ private struct ArrowKeyedDecoding<Key: CodingKey>: KeyedDecodingContainerProtoco
}
func decode(_ type: Int.Type, forKey key: Key) throws -> Int {
- throw ArrowError.invalid(
- "Int type is not supported (please use Int8, Int16, Int32 or Int64)")
+ let col = try self.decoder.getCol(key.stringValue)
+ return try self.decoder.decodeInt(from: col, keyDescription: key.stringValue)
}
func decode(_ type: Int8.Type, forKey key: Key) throws -> Int8 {
@@ -244,8 +342,8 @@ private struct ArrowKeyedDecoding<Key: CodingKey>: KeyedDecodingContainerProtoco
}
func decode(_ type: UInt.Type, forKey key: Key) throws -> UInt {
- throw ArrowError.invalid(
- "UInt type is not supported (please use UInt8, UInt16, UInt32 or UInt64)")
+ let col = try self.decoder.getCol(key.stringValue)
+ return try self.decoder.decodeUInt(from: col, keyDescription: key.stringValue)
}
func decode(_ type: UInt8.Type, forKey key: Key) throws -> UInt8 {
@@ -268,8 +366,26 @@ private struct ArrowKeyedDecoding<Key: CodingKey>: KeyedDecodingContainerProtoco
return try self.decoder.doDecode(key)!
}
+ func decode(_ type: Date.Type, forKey key: Key) throws -> Date {
+ let col = try self.decoder.getCol(key.stringValue)
+ return try self.decoder.decodeDate(from: col, keyDescription: key.stringValue)
+ }
+
+ func decode(_ type: TimestampNanos.Type, forKey key: Key) throws -> TimestampNanos {
+ let col = try self.decoder.getCol(key.stringValue)
+ return try self.decoder.decodeTimestampNanos(from: col, keyDescription: key.stringValue)
+ }
+
func decode<T>(_ type: T.Type, forKey key: Key) throws -> T where T: Decodable {
- if ArrowArrayBuilders.isValidBuilderType(type) || type == Date.self || type == Decimal.self {
+ if type == Int.self {
+ return try decode(Int.self, forKey: key) as! T
+ } else if type == UInt.self {
+ return try decode(UInt.self, forKey: key) as! T
+ } else if type == Date.self {
+ return try decode(Date.self, forKey: key) as! T
+ } else if type == TimestampNanos.self {
+ return try decode(TimestampNanos.self, forKey: key) as! T
+ } else if ArrowArrayBuilders.isValidBuilderType(type) || type == Decimal.self {
return try self.decoder.doDecode(key)!
} else {
throw ArrowError.invalid("Type \(type) is currently not supported")
@@ -330,8 +446,8 @@ private struct ArrowSingleValueDecoding: SingleValueDecodingContainer {
}
func decode(_ type: Int.Type) throws -> Int {
- throw ArrowError.invalid(
- "Int type is not supported (please use Int8, Int16, Int32 or Int64)")
+ let col = try self.decoder.getCol(self.decoder.singleRBCol)
+ return try self.decoder.decodeInt(from: col, keyDescription: "column \(self.decoder.singleRBCol)")
}
func decode(_ type: Int8.Type) throws -> Int8 {
@@ -351,8 +467,8 @@ private struct ArrowSingleValueDecoding: SingleValueDecodingContainer {
}
func decode(_ type: UInt.Type) throws -> UInt {
- throw ArrowError.invalid(
- "UInt type is not supported (please use UInt8, UInt16, UInt32 or UInt64)")
+ let col = try self.decoder.getCol(self.decoder.singleRBCol)
+ return try self.decoder.decodeUInt(from: col, keyDescription: "column \(self.decoder.singleRBCol)")
}
func decode(_ type: UInt8.Type) throws -> UInt8 {
@@ -375,8 +491,27 @@ private struct ArrowSingleValueDecoding: SingleValueDecodingContainer {
return try self.decoder.doDecode(self.decoder.singleRBCol)!
}
+ func decode(_ type: Date.Type) throws -> Date {
+ let col = try self.decoder.getCol(self.decoder.singleRBCol)
+ return try self.decoder.decodeDate(from: col, keyDescription: "column \(self.decoder.singleRBCol)")
+ }
+
+ func decode(_ type: TimestampNanos.Type) throws -> TimestampNanos {
+ let col = try self.decoder.getCol(self.decoder.singleRBCol)
+ return try self.decoder.decodeTimestampNanos(
+ from: col, keyDescription: "column \(self.decoder.singleRBCol)")
+ }
+
func decode<T>(_ type: T.Type) throws -> T where T: Decodable {
- if ArrowArrayBuilders.isValidBuilderType(type) || type == Date.self || type == Decimal.self {
+ if type == Int.self {
+ return try decode(Int.self) as! T
+ } else if type == UInt.self {
+ return try decode(UInt.self) as! T
+ } else if type == Date.self {
+ return try decode(Date.self) as! T
+ } else if type == TimestampNanos.self {
+ return try decode(TimestampNanos.self) as! T
+ } else if ArrowArrayBuilders.isValidBuilderType(type) || type == Decimal.self {
return try self.decoder.doDecode(self.decoder.singleRBCol)!
} else {
throw ArrowError.invalid("Type \(type) is currently not supported")
diff --git a/Sources/SparkConnect/ArrowEncoder.swift b/Sources/SparkConnect/ArrowEncoder.swift
index 1f85941..792b390 100644
--- a/Sources/SparkConnect/ArrowEncoder.swift
+++ b/Sources/SparkConnect/ArrowEncoder.swift
@@ -67,6 +67,16 @@ public class ArrowEncoder: Encoder {
if ArrowArrayBuilders.isValidBuilderType(T.self) {
let builders = ["col0": try ArrowArrayBuilders.loadBuilder(T.self)]
return ArrowEncoder(builders, byIndex: ["col0"])
+ } else if T.self == Int.self || T.self == Int?.self {
+ let builders = [
+ "col0": try ArrowArrayBuilders.loadNumberArrayBuilder() as NumberArrayBuilder<Int64>
+ ]
+ return ArrowEncoder(builders, byIndex: ["col0"])
+ } else if T.self == UInt.self || T.self == UInt?.self {
+ let builders = [
+ "col0": try ArrowArrayBuilders.loadNumberArrayBuilder() as NumberArrayBuilder<UInt64>
+ ]
+ return ArrowEncoder(builders, byIndex: ["col0"])
} else {
let encoder = ArrowEncoder()
if data is [AnyHashable: Any] {
@@ -236,13 +246,11 @@ private struct ArrowKeyedEncoding<Key: CodingKey>: KeyedEncodingContainerProtoco
}
mutating func encode(_ value: Int, forKey key: Key) throws {
- throw ArrowError.invalid(
- "Int type is not supported (please use Int8, Int16, Int32 or Int64)")
+ try encoder.doEncode(Int64(value), key: key)
}
mutating func encodeIfPresent(_ value: Int?, forKey key: Key) throws {
- throw ArrowError.invalid(
- "Int type is not supported (please use Int8, Int16, Int32 or Int64)")
+ try doEncodeIf(value.map(Int64.init), forKey: key)
}
mutating func encode(_ value: Int8, forKey key: Key) throws {
@@ -278,13 +286,11 @@ private struct ArrowKeyedEncoding<Key: CodingKey>: KeyedEncodingContainerProtoco
}
mutating func encode(_ value: UInt, forKey key: Key) throws {
- throw ArrowError.invalid(
- "UInt type is not supported (please use UInt8, UInt16, UInt32 or UInt64)")
+ try encoder.doEncode(UInt64(value), key: key)
}
mutating func encodeIfPresent(_ value: UInt?, forKey key: Key) throws {
- throw ArrowError.invalid(
- "UInt type is not supported (please use UInt8, UInt16, UInt32 or UInt64)")
+ try doEncodeIf(value.map(UInt64.init), forKey: key)
}
mutating func encode(_ value: UInt8, forKey key: Key) throws {
@@ -320,7 +326,11 @@ private struct ArrowKeyedEncoding<Key: CodingKey>: KeyedEncodingContainerProtoco
}
mutating func encode<T: Encodable>(_ value: T, forKey key: Key) throws {
- if ArrowArrayBuilders.isValidBuilderType(T.self) {
+ if let intVal = value as? Int {
+ try encoder.doEncode(Int64(intVal), key: key)
+ } else if let uintVal = value as? UInt {
+ try encoder.doEncode(UInt64(uintVal), key: key)
+ } else if ArrowArrayBuilders.isValidBuilderType(T.self) {
try encoder.doEncode(value, key: key)
} else {
throw ArrowError.invalid("Type \(T.self) is currently not supported")
@@ -328,7 +338,11 @@ private struct ArrowKeyedEncoding<Key: CodingKey>: KeyedEncodingContainerProtoco
}
mutating func encodeIfPresent<T>(_ value: T?, forKey key: Self.Key) throws where T: Encodable {
- if ArrowArrayBuilders.isValidBuilderType(T?.self) {
+ if let intVal = value as? Int {
+ try doEncodeIf(Int64(intVal), forKey: key)
+ } else if let uintVal = value as? UInt {
+ try doEncodeIf(UInt64(uintVal), forKey: key)
+ } else if ArrowArrayBuilders.isValidBuilderType(T?.self) {
try doEncodeIf(value, forKey: key)
} else {
throw ArrowError.invalid("Type \(T.self) is currently not supported")
@@ -402,12 +416,17 @@ private struct ArrowUnkeyedEncoding: UnkeyedEncodingContainer {
}
mutating func encode<T>(_ value: T) throws where T: Encodable {
- let type = T.self
- if ArrowArrayBuilders.isValidBuilderType(type) {
+ if let intVal = value as? Int {
+ defer { increment() }
+ return try self.encoder.doEncode(Int64(intVal), keyIndex: self.currentIndex)
+ } else if let uintVal = value as? UInt {
+ defer { increment() }
+ return try self.encoder.doEncode(UInt64(uintVal), keyIndex: self.currentIndex)
+ } else if ArrowArrayBuilders.isValidBuilderType(T.self) {
defer { increment() }
return try self.encoder.doEncode(value, keyIndex: self.currentIndex)
} else {
- throw ArrowError.invalid("Type \(type) is currently not supported")
+ throw ArrowError.invalid("Type \(T.self) is currently not supported")
}
}
@@ -454,7 +473,11 @@ private struct ArrowSingleValueEncoding: SingleValueEncodingContainer {
}
mutating func encode<T: Encodable>(_ value: T) throws {
- if ArrowArrayBuilders.isValidBuilderType(T.self) {
+ if let intVal = value as? Int {
+ return try self.encoder.doEncode(Int64(intVal), keyIndex: 0)
+ } else if let uintVal = value as? UInt {
+ return try self.encoder.doEncode(UInt64(uintVal), keyIndex: 0)
+ } else if ArrowArrayBuilders.isValidBuilderType(T.self) {
return try self.encoder.doEncode(value, keyIndex: 0)
} else {
throw ArrowError.invalid("Type \(T.self) is currently not supported")
diff --git a/Sources/SparkConnect/DataFrame+Actions.swift b/Sources/SparkConnect/DataFrame+Actions.swift
index 136471b..86a2443 100644
--- a/Sources/SparkConnect/DataFrame+Actions.swift
+++ b/Sources/SparkConnect/DataFrame+Actions.swift
@@ -194,6 +194,31 @@ extension DataFrame {
return result
}
+ /// Execute the plan and return the result decoded into an array of `T`.
+ ///
+ /// ```swift
+ /// struct Person: Codable, Sendable, Equatable {
+ /// let name: String
+ /// let age: Int
+ /// }
+ ///
+ /// let people: [Person] = try await df.collect(as: Person.self)
+ /// ```
+ ///
+ /// - Parameter type: The `Decodable` type to decode each row into. Defaults to `T.self`.
+ /// - Returns: An array of decoded instances of `T`.
+ public func collect<T: Decodable>(as type: T.Type = T.self) async throws -> [T] {
+ try await execute()
+
+ var result: [T] = []
+ for batch in self.batches {
+ let decoder = ArrowDecoder(batch)
+ let decoded = try decoder.decode(type)
+ result.append(contentsOf: decoded)
+ }
+ return result
+ }
+
// MARK: - Count
/// Return the total number of rows.
diff --git a/Sources/SparkConnect/DataType.swift b/Sources/SparkConnect/DataType.swift
index 7d6747e..6c6d8c5 100644
--- a/Sources/SparkConnect/DataType.swift
+++ b/Sources/SparkConnect/DataType.swift
@@ -321,6 +321,45 @@ extension DataType {
return proto
}
+ /// Create a public ``DataType`` from an Apache Arrow ``ArrowType``.
+ init(_ arrowType: ArrowType) throws {
+ switch arrowType.id {
+ case .boolean:
+ self = .boolean
+ case .int8:
+ self = .byte
+ case .int16, .uint8:
+ self = .short
+ case .int32, .uint16:
+ self = .integer
+ case .int64, .uint32:
+ self = .long
+ case .uint64:
+ self = .decimal(precision: 20, scale: 0)
+ case .float:
+ self = .float
+ case .double:
+ self = .double
+ case .decimal128(let precision, let scale):
+ self = .decimal(precision: precision, scale: scale)
+ case .string:
+ self = .string
+ case .binary:
+ self = .binary
+ case .date32, .date64:
+ self = .date
+ case .timestamp:
+ self = .timestamp
+ case .time32:
+ self = .time(precision: 3)
+ case .time64:
+ let unit = (arrowType as? ArrowTypeTime64)?.unit
+ self = .time(precision: unit == .microseconds ? 6 : 9)
+ default:
+ throw SparkConnectError.InvalidType
+ }
+ }
+
/// Create a public ``DataType`` from the `Spark Connect` protobuf representation.
init(_ proto: ProtoDataType) throws {
switch proto.kind {
diff --git a/Sources/SparkConnect/SparkSession.swift b/Sources/SparkConnect/SparkSession.swift
index c5cb6b0..ad8ca93 100644
--- a/Sources/SparkConnect/SparkSession.swift
+++ b/Sources/SparkConnect/SparkSession.swift
@@ -164,16 +164,7 @@ public actor SparkSession {
throw SparkConnectError.InvalidType
}
let ipcStream = try ConvertToArrow.toArrowIPCStream(data, structType)
- let threshold = Int(
- try await client.getConf("spark.sql.session.localRelationCacheThreshold"))!
- if ipcStream.count >= threshold {
- let hash = try await client.cacheLocalRelation(ipcStream, schema)
- return await DataFrame(spark: self, plan: client.getCachedLocalRelation(hash))
- }
- guard ipcStream.count <= Self.maxLocalRelationSize else {
- throw SparkConnectError.LocalRelationTooLarge
- }
- return await DataFrame(spark: self, plan: client.getLocalRelation(ipcStream, schema))
+ return try await createDataFrame(ipcStream: ipcStream, schema: schema)
}
/// Creates a ``DataFrame`` from the given local data with the specified ``StructType`` schema.
@@ -187,6 +178,87 @@ public actor SparkSession {
return try await createDataFrame(data, schema.toDDL)
}
+ /// Creates a ``DataFrame`` from the given array of `Encodable` instances.
+ ///
+ /// The schema is automatically inferred from the properties of `T` using `ArrowEncoder`.
+ ///
+ /// ```swift
+ /// struct Person: Codable, Sendable, Equatable {
+ /// let name: String
+ /// let age: Int
+ /// }
+ ///
+ /// let data = [Person(name: "Alice", age: 20), Person(name: "Bob", age: 25)]
+ /// let df = try await spark.createDataFrame(data)
+ /// ```
+ ///
+ /// - Parameter data: A non-empty array of `Encodable` items.
+ /// - Returns: A ``DataFrame`` instance.
+ /// - Throws: ``SparkConnectError/invalidArgument(_:)`` if `data` is empty, or if encoding fails.
+ public func createDataFrame<T: Encodable>(_ data: [T]) async throws -> DataFrame {
+ guard !data.isEmpty else {
+ throw SparkConnectError.invalidArgument(
+ SparkConnectError.Details(message: "Cannot infer schema from empty data"))
+ }
+ guard let batch = try ArrowEncoder.encode(data) else {
+ throw SparkConnectError.InvalidArrowData
+ }
+ let schema = try StructType(batch.schema).toDDL
+ return try await createDataFrame(batch: batch, schema: schema)
+ }
+
+ /// Creates a ``DataFrame`` from the given array of `Encodable` instances with the specified schema string.
+ /// - Parameters:
+ /// - data: An array of `Encodable` items.
+ /// - schema: A DDL-formatted schema string, e.g. `id INT, name STRING`.
+ /// - Returns: A ``DataFrame`` instance.
+ @_disfavoredOverload
+ public func createDataFrame<T: Encodable>(_ data: [T], _ schema: String) async throws -> DataFrame {
+ if data.isEmpty {
+ return try await createDataFrame([], schema)
+ }
+ guard let batch = try ArrowEncoder.encode(data) else {
+ throw SparkConnectError.InvalidArrowData
+ }
+ return try await createDataFrame(batch: batch, schema: schema)
+ }
+
+ /// Creates a ``DataFrame`` from the given array of `Encodable` instances with the specified ``StructType`` schema.
+ /// - Parameters:
+ /// - data: An array of `Encodable` items.
+ /// - schema: A ``StructType`` schema.
+ /// - Returns: A ``DataFrame`` instance.
+ @_disfavoredOverload
+ public func createDataFrame<T: Encodable>(_ data: [T], _ schema: StructType) async throws
+ -> DataFrame
+ {
+ return try await createDataFrame(data, schema.toDDL)
+ }
+
+ private func createDataFrame(batch: RecordBatch, schema: String) async throws -> DataFrame {
+ switch ArrowWriter().writeStreaming(
+ ArrowWriter.Info(.recordbatch, schema: batch.schema, batches: [batch]))
+ {
+ case .success(let ipcStream):
+ return try await createDataFrame(ipcStream: ipcStream, schema: schema)
+ case .failure:
+ throw SparkConnectError.InvalidArrowData
+ }
+ }
+
+ private func createDataFrame(ipcStream: Data, schema: String) async throws -> DataFrame {
+ let threshold = Int(
+ try await client.getConf("spark.sql.session.localRelationCacheThreshold"))!
+ if ipcStream.count >= threshold {
+ let hash = try await client.cacheLocalRelation(ipcStream, schema)
+ return await DataFrame(spark: self, plan: client.getCachedLocalRelation(hash))
+ }
+ guard ipcStream.count <= Self.maxLocalRelationSize else {
+ throw SparkConnectError.LocalRelationTooLarge
+ }
+ return await DataFrame(spark: self, plan: client.getLocalRelation(ipcStream, schema))
+ }
+
/// Parse a DDL-formatted schema string into a ``StructType`` using the server-side parser.
///
/// ```swift
diff --git a/Sources/SparkConnect/StructType.swift b/Sources/SparkConnect/StructType.swift
index 5550290..243755b 100644
--- a/Sources/SparkConnect/StructType.swift
+++ b/Sources/SparkConnect/StructType.swift
@@ -68,6 +68,16 @@ public struct StructType: Sendable, Equatable {
self.fields = fields
}
+ /// Create a ``StructType`` from an Apache Arrow ``ArrowSchema``.
+ init(_ arrowSchema: ArrowSchema) throws {
+ var fields = [StructField]()
+ for field in arrowSchema.fields {
+ let dataType = try DataType(field.type)
+ fields.append(StructField(name: field.name, dataType: dataType, nullable: field.isNullable))
+ }
+ self.init(fields: fields)
+ }
+
/// The names of all fields in order.
public var fieldNames: [String] {
fields.map { $0.name }
diff --git a/Tests/SparkConnectTests/CreateDataFrameTests.swift b/Tests/SparkConnectTests/CreateDataFrameTests.swift
index 7122c8a..649f68a 100644
--- a/Tests/SparkConnectTests/CreateDataFrameTests.swift
+++ b/Tests/SparkConnectTests/CreateDataFrameTests.swift
@@ -223,4 +223,80 @@ struct CreateDataFrameTests {
#expect(try await df2.count() == 10000)
await spark.stop()
}
+
+ struct Person: Codable, Sendable, Equatable {
+ let name: String
+ let age: Int
+ }
+
+ struct UserProfile: Codable, Sendable, Equatable {
+ let id: Int
+ let nickname: String?
+ let score: Double?
+ }
+
+ @Test
+ func createDataFrameWithEncodable() async throws {
+ let spark = try await SparkSession.builder.getOrCreate()
+ let data = [
+ Person(name: "Alice", age: 20),
+ Person(name: "Bob", age: 25),
+ ]
+ let df = try await spark.createDataFrame(data)
+ #expect(try await df.columns == ["name", "age"])
+ #expect(try await df.count() == 2)
+ let collected: [Person] = try await df.collect(as: Person.self)
+ #expect(collected == data)
+ await spark.stop()
+ }
+
+ @Test
+ func createDataFrameWithEncodableOptionals() async throws {
+ let spark = try await SparkSession.builder.getOrCreate()
+ let data = [
+ UserProfile(id: 1, nickname: nil, score: 98.5),
+ UserProfile(id: 2, nickname: "spark_user", score: nil),
+ ]
+ let df = try await spark.createDataFrame(data)
+ #expect(try await df.count() == 2)
+ let collected: [UserProfile] = try await df.collect(as: UserProfile.self)
+ #expect(collected == data)
+ await spark.stop()
+ }
+
+ @Test
+ func createDataFrameWithEncodableEmpty() async throws {
+ let spark = try await SparkSession.builder.getOrCreate()
+ let empty: [Person] = []
+ await #expect(throws: SparkConnectError.self) {
+ try await spark.createDataFrame(empty)
+ }
+
+ // With explicit schema, empty encodable array works:
+ let df = try await spark.createDataFrame(empty, "name STRING, age INT")
+ #expect(try await df.count() == 0)
+ let collected: [Person] = try await df.collect(as: Person.self)
+ #expect(collected == [])
+ await spark.stop()
+ }
+
+ @Test
+ func createDataFrameWithEncodablePrimitives() async throws {
+ let spark = try await SparkSession.builder.getOrCreate()
+ let numbers = [10, 20, 30]
+ let df = try await spark.createDataFrame(numbers)
+ #expect(try await df.count() == 3)
+ let collected: [Int] = try await df.collect(as: Int.self)
+ #expect(collected == numbers)
+ await spark.stop()
+ }
+
+ @Test
+ func collectAsWithSparkQuery() async throws {
+ let spark = try await SparkSession.builder.getOrCreate()
+ let df = try await spark.sql("SELECT 'Alice' AS name, 30 AS age")
+ let people: [Person] = try await df.collect(as: Person.self)
+ #expect(people == [Person(name: "Alice", age: 30)])
+ await spark.stop()
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org