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 3a3a0f3 [SPARK-58648] Support StructType in schema, createDataFrame,
and to APIs
3a3a0f3 is described below
commit 3a3a0f39bc577bfb84e353dd2d02cbcbfaa17fb2
Author: Dongjoon Hyun <[email protected]>
AuthorDate: Fri Aug 7 12:07:33 2026 -0700
[SPARK-58648] Support StructType in schema, createDataFrame, and to APIs
### What changes were proposed in this pull request?
This PR adds `StructType` overloads to the four schema-taking APIs which
previously accepted only DDL strings.
- `DataFrameReader.schema(_ schema: StructType)`
- `DataStreamReader.schema(_ schema: StructType)`
- `SparkSession.createDataFrame(_ data: [[Sendable?]], _ schema:
StructType)`
- `DataFrame.to(_ schema: StructType)`
Following the Scala client implementation, the reader and `createDataFrame`
overloads convert the schema via the newly added public
`StructType.toDDL`/`StructField.toDDL` (Catalyst `toDDL` format with
backtick-quoting and `NOT NULL`) and reuse the existing DDL-string paths, while
`DataFrame.to` converts the model directly to the protobuf `DataType` (like
`DataTypeProtoConverter.toConnectProtoType`) and requires no server round-trip.
```swift
let schema = StructType(fields: [
StructField(name: "id", dataType: .integer, nullable: false),
StructField(name: "name", dataType: .string),
])
let df = try await spark.createDataFrame([[1, "Alice"], [2, nil]], schema)
```
### Why are the changes needed?
For feature parity with the Scala/PySpark clients where all schema-taking
APIs accept a `StructType` object, so users can build and pass schemas
programmatically instead of assembling DDL strings by hand.
### Does this PR introduce _any_ user-facing change?
No. This adds new APIs (`StructType` overloads and `toDDL`) without
changing existing behavior.
### 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: Claude Fable 5
Closes #488 from dongjoon-hyun/SPARK-58648.
Authored-by: Dongjoon Hyun <[email protected]>
Signed-off-by: Dongjoon Hyun <[email protected]>
---
.../SparkConnect/DataFrame+Transformations.swift | 9 ++
Sources/SparkConnect/DataFrameReader.swift | 10 ++
Sources/SparkConnect/DataStreamReader.swift | 10 ++
Sources/SparkConnect/DataType.swift | 126 +++++++++++++++++++++
Sources/SparkConnect/SparkSession.swift | 11 ++
Sources/SparkConnect/StructType.swift | 54 +++++++++
Sources/SparkConnect/TypeAliases.swift | 1 +
Tests/SparkConnectTests/CreateDataFrameTests.swift | 19 ++++
Tests/SparkConnectTests/DataFrameReaderTests.swift | 13 +++
Tests/SparkConnectTests/DataFrameTests.swift | 25 ++++
Tests/SparkConnectTests/DataStreamTests.swift | 14 +++
Tests/SparkConnectTests/StructTypeTests.swift | 25 ++++
12 files changed, 317 insertions(+)
diff --git a/Sources/SparkConnect/DataFrame+Transformations.swift
b/Sources/SparkConnect/DataFrame+Transformations.swift
index ccb83c9..492f9b5 100644
--- a/Sources/SparkConnect/DataFrame+Transformations.swift
+++ b/Sources/SparkConnect/DataFrame+Transformations.swift
@@ -69,6 +69,15 @@ extension DataFrame {
spark: self.spark, plan: SparkConnectClient.getToSchema(self.plan.root,
dataType))
}
+ /// Returns a new DataFrame where each row is reconciled to match the
specified schema.
+ /// - Parameter schema: The given ``StructType`` schema.
+ /// - Returns: A ``DataFrame`` with the given schema.
+ public func to(_ schema: StructType) -> DataFrame {
+ return DataFrame(
+ spark: self.spark,
+ plan: SparkConnectClient.getToSchema(self.plan.root,
DataType.struct(schema).toProtoDataType))
+ }
+
/// Returns the content of the Dataset as a Dataset of JSON strings.
/// - Returns: A ``DataFrame`` with a single string column whose content is
JSON.
public func toJSON() -> DataFrame {
diff --git a/Sources/SparkConnect/DataFrameReader.swift
b/Sources/SparkConnect/DataFrameReader.swift
index 28d07ac..245ae2e 100644
--- a/Sources/SparkConnect/DataFrameReader.swift
+++ b/Sources/SparkConnect/DataFrameReader.swift
@@ -132,6 +132,16 @@ public actor DataFrameReader: Sendable {
return self
}
+ /// Specifies the input schema. Some data sources (e.g. JSON) can infer the
input schema
+ /// automatically from data. By specifying the schema here, the underlying
data source can skip
+ /// the schema inference step, and thus speed up data loading.
+ /// - Parameter schema: A ``StructType`` schema.
+ /// - Returns: A ``DataFrameReader``.
+ @discardableResult
+ public func schema(_ schema: StructType) async throws -> DataFrameReader {
+ return try await self.schema(schema.toDDL)
+ }
+
/// Loads input in as a ``DataFrame``, for data sources that don't require a
path (e.g. external
/// key-value stores).
/// - Returns: A ``DataFrame``.
diff --git a/Sources/SparkConnect/DataStreamReader.swift
b/Sources/SparkConnect/DataStreamReader.swift
index e1be7dd..931ce99 100644
--- a/Sources/SparkConnect/DataStreamReader.swift
+++ b/Sources/SparkConnect/DataStreamReader.swift
@@ -59,6 +59,16 @@ public actor DataStreamReader: Sendable {
return self
}
+ /// Specifies the input schema. Some data sources (e.g. JSON) can infer the
input schema
+ /// automatically from data. By specifying the schema here, the underlying
data source can skip
+ /// the schema inference step, and thus speed up data loading.
+ /// - Parameter schema: A ``StructType`` schema.
+ /// - Returns: A ``DataStreamReader``.
+ @discardableResult
+ public func schema(_ schema: StructType) async throws -> DataStreamReader {
+ return try await self.schema(schema.toDDL)
+ }
+
/// Adds an input option for the underlying data source.
/// - Parameters:
/// - key: A key string.
diff --git a/Sources/SparkConnect/DataType.swift
b/Sources/SparkConnect/DataType.swift
index 7fe7f54..5611204 100644
--- a/Sources/SparkConnect/DataType.swift
+++ b/Sources/SparkConnect/DataType.swift
@@ -181,6 +181,132 @@ public indirect enum DataType: Sendable, Equatable {
}
extension DataType {
+ /// A type string usable in DDL like Spark SQL's `DataType.sql`, e.g. `INT`
or
+ /// `STRUCT<x: INT NOT NULL>`. Unlike ``simpleString``, struct field names
are quoted
+ /// if needed and their `NOT NULL` constraints are kept.
+ var sql: String {
+ switch self {
+ case .array(let elementType, _):
+ "ARRAY<\(elementType.sql)>"
+ case .map(let keyType, let valueType, _):
+ "MAP<\(keyType.sql), \(valueType.sql)>"
+ case .struct(let structType):
+ structType.sql
+ case .udt(let udt):
+ udt.sqlType?.sql ?? "UDT"
+ case .unparsed(let dataTypeString):
+ dataTypeString
+ default:
+ simpleString.uppercased()
+ }
+ }
+
+ /// Convert to the `Spark Connect` protobuf representation.
+ var toProtoDataType: ProtoDataType {
+ var proto = ProtoDataType()
+ switch self {
+ case .null:
+ proto.null = ProtoDataType.NULL()
+ case .binary:
+ proto.binary = ProtoDataType.Binary()
+ case .boolean:
+ proto.boolean = ProtoDataType.Boolean()
+ case .byte:
+ proto.byte = ProtoDataType.Byte()
+ case .short:
+ proto.short = ProtoDataType.Short()
+ case .integer:
+ proto.integer = ProtoDataType.Integer()
+ case .long:
+ proto.long = ProtoDataType.Long()
+ case .float:
+ proto.float = ProtoDataType.FloatMessage()
+ case .double:
+ proto.double = ProtoDataType.DoubleMessage()
+ case .decimal(let precision, let scale):
+ var decimal = ProtoDataType.Decimal()
+ decimal.precision = precision
+ decimal.scale = scale
+ proto.decimal = decimal
+ case .string:
+ proto.string = ProtoDataType.StringMessage()
+ case .char(let length):
+ var char = ProtoDataType.Char()
+ char.length = length
+ proto.char = char
+ case .varchar(let length):
+ var varChar = ProtoDataType.VarChar()
+ varChar.length = length
+ proto.varChar = varChar
+ case .date:
+ proto.date = ProtoDataType.Date()
+ case .timestamp:
+ proto.timestamp = ProtoDataType.Timestamp()
+ case .timestampNtz:
+ proto.timestampNtz = ProtoDataType.TimestampNTZ()
+ case .time(let precision):
+ var time = ProtoDataType.Time()
+ time.precision = precision
+ proto.time = time
+ case .calendarInterval:
+ proto.calendarInterval = ProtoDataType.CalendarInterval()
+ case .yearMonthInterval(let startField, let endField):
+ var interval = ProtoDataType.YearMonthInterval()
+ interval.startField = startField.rawValue
+ interval.endField = endField.rawValue
+ proto.yearMonthInterval = interval
+ case .dayTimeInterval(let startField, let endField):
+ var interval = ProtoDataType.DayTimeInterval()
+ interval.startField = startField.rawValue
+ interval.endField = endField.rawValue
+ proto.dayTimeInterval = interval
+ case .array(let elementType, let containsNull):
+ var array = ProtoDataType.Array()
+ array.elementType = elementType.toProtoDataType
+ array.containsNull = containsNull
+ proto.array = array
+ case .map(let keyType, let valueType, let valueContainsNull):
+ var map = ProtoDataType.Map()
+ map.keyType = keyType.toProtoDataType
+ map.valueType = valueType.toProtoDataType
+ map.valueContainsNull = valueContainsNull
+ proto.map = map
+ case .struct(let structType):
+ proto.struct = structType.toProtoStructType
+ case .variant:
+ proto.variant = ProtoDataType.Variant()
+ case .geometry(let srid):
+ var geometry = ProtoDataType.Geometry()
+ geometry.srid = srid
+ proto.geometry = geometry
+ case .geography(let srid):
+ var geography = ProtoDataType.Geography()
+ geography.srid = srid
+ proto.geography = geography
+ case .udt(let udt):
+ var udtProto = ProtoDataType.UDT()
+ udtProto.type = "udt"
+ if let jvmClass = udt.jvmClass {
+ udtProto.jvmClass = jvmClass
+ }
+ if let pythonClass = udt.pythonClass {
+ udtProto.pythonClass = pythonClass
+ }
+ if let serializedPythonClass = udt.serializedPythonClass {
+ udtProto.serializedPythonClass = serializedPythonClass
+ }
+ if let sqlType = udt.sqlType {
+ udtProto.sqlType = sqlType.toProtoDataType
+ }
+ proto.udt = udtProto
+ case .unparsed(let dataTypeString):
+ var unparsed = ProtoDataType.Unparsed()
+ unparsed.dataTypeString = dataTypeString
+ proto.unparsed = unparsed
+ }
+ return proto
+ }
+
/// 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 907dd91..27e1f38 100644
--- a/Sources/SparkConnect/SparkSession.swift
+++ b/Sources/SparkConnect/SparkSession.swift
@@ -175,6 +175,17 @@ public actor SparkSession {
return await DataFrame(spark: self, plan:
client.getLocalRelation(ipcStream, schema))
}
+ /// Creates a ``DataFrame`` from the given local data with the specified
``StructType`` schema.
+ /// This is equivalent to calling `createDataFrame(_:_:)` with the DDL
string of the given schema.
+ /// - Parameters:
+ /// - data: An array of rows whose values are ordered like the schema
fields.
+ /// - schema: A ``StructType`` schema.
+ /// - Returns: A ``DataFrame`` instance.
+ public func createDataFrame(_ data: [[Sendable?]], _ schema: StructType)
async throws -> DataFrame
+ {
+ return try await createDataFrame(data, schema.toDDL)
+ }
+
/// Create a ``DataFrame`` with a single `Int64` column name `id`,
containing elements in a
/// range from 0 to `end` (exclusive) with step value 1.
///
diff --git a/Sources/SparkConnect/StructType.swift
b/Sources/SparkConnect/StructType.swift
index 3896230..5550290 100644
--- a/Sources/SparkConnect/StructType.swift
+++ b/Sources/SparkConnect/StructType.swift
@@ -30,6 +30,32 @@ public struct StructField: Sendable, Equatable {
self.nullable = nullable
self.metadata = metadata
}
+
+ /// A string containing this field in DDL format like Spark SQL's
`StructField.toDDL`,
+ /// e.g. `` `eventId` INT NOT NULL ``.
+ public var toDDL: String {
+ "\(quoteIfNeeded(name)) \(dataType.sql)\(nullable ? "" : " NOT NULL")"
+ }
+
+ /// A string usable inside a struct type string like Spark SQL's
`StructField.sql`,
+ /// e.g. `` `eventId`: INT NOT NULL ``.
+ var sql: String {
+ "\(quoteIfNeeded(name)): \(dataType.sql)\(nullable ? "" : " NOT NULL")"
+ }
+}
+
+/// Quote the given name with backticks like Spark SQL's
`QuotingUtils.quoteIfNeeded`
+/// unless it is a valid identifier matching `[a-zA-Z_][a-zA-Z0-9_]*`.
+private func quoteIfNeeded(_ name: String) -> String {
+ let isValidIdentifier =
+ !name.isEmpty
+ && name.utf8.enumerated().allSatisfy { (index, byte) in
+ byte == UInt8(ascii: "_")
+ || (UInt8(ascii: "a")...UInt8(ascii: "z")).contains(byte)
+ || (UInt8(ascii: "A")...UInt8(ascii: "Z")).contains(byte)
+ || (index > 0 && (UInt8(ascii: "0")...UInt8(ascii:
"9")).contains(byte))
+ }
+ return isValidIdentifier ? name : "`\(name.replacing("`", with: "``"))`"
}
/// A struct type holding an ordered collection of ``StructField``s, mirroring
Spark SQL's
@@ -57,6 +83,18 @@ public struct StructType: Sendable, Equatable {
public var simpleString: String {
"struct<\(fields.map { "\($0.name):\($0.dataType.simpleString)"
}.joined(separator: ","))>"
}
+
+ /// A string containing this schema in DDL format like Spark SQL's
`StructType.toDDL`,
+ /// e.g. `id BIGINT NOT NULL,name STRING`.
+ public var toDDL: String {
+ fields.map { $0.toDDL }.joined(separator: ",")
+ }
+
+ /// A type string usable in DDL like Spark SQL's `StructType.sql`,
+ /// e.g. `STRUCT<id: BIGINT NOT NULL, name: STRING>`.
+ var sql: String {
+ "STRUCT<\(fields.map { $0.sql }.joined(separator: ", "))>"
+ }
}
extension StructType: RandomAccessCollection {
@@ -69,6 +107,22 @@ extension StructType: RandomAccessCollection {
}
extension StructType {
+ /// Convert to the `Spark Connect` protobuf representation.
+ var toProtoStructType: ProtoStructType {
+ var proto = ProtoStructType()
+ proto.fields = fields.map { field in
+ var protoField = ProtoStructField()
+ protoField.name = field.name
+ protoField.dataType = field.dataType.toProtoDataType
+ protoField.nullable = field.nullable
+ if let metadata = field.metadata {
+ protoField.metadata = metadata
+ }
+ return protoField
+ }
+ return proto
+ }
+
/// Create a public ``StructType`` from the `Spark Connect` protobuf
representation.
init(_ proto: ProtoStructType) throws {
self.fields = try proto.fields.map {
diff --git a/Sources/SparkConnect/TypeAliases.swift
b/Sources/SparkConnect/TypeAliases.swift
index 43c8a00..788e03e 100644
--- a/Sources/SparkConnect/TypeAliases.swift
+++ b/Sources/SparkConnect/TypeAliases.swift
@@ -52,6 +52,7 @@ typealias Parse = Spark_Connect_Parse
typealias Plan = Spark_Connect_Plan
typealias Project = Spark_Connect_Project
typealias ProtoDataType = Spark_Connect_DataType
+typealias ProtoStructField = Spark_Connect_DataType.StructField
typealias ProtoStructType = Spark_Connect_DataType.Struct
typealias Range = Spark_Connect_Range
typealias Read = Spark_Connect_Read
diff --git a/Tests/SparkConnectTests/CreateDataFrameTests.swift
b/Tests/SparkConnectTests/CreateDataFrameTests.swift
index 11947b8..4858713 100644
--- a/Tests/SparkConnectTests/CreateDataFrameTests.swift
+++ b/Tests/SparkConnectTests/CreateDataFrameTests.swift
@@ -40,6 +40,25 @@ struct CreateDataFrameTests {
await spark.stop()
}
+ @Test
+ func createDataFrameWithStructType() async throws {
+ let spark = try await SparkSession.builder.getOrCreate()
+ let schema = StructType(fields: [
+ StructField(name: "id", dataType: .integer, nullable: false),
+ StructField(name: "name", dataType: .string),
+ ])
+ let df = try await spark.createDataFrame([[1, "Alice"], [2, nil]], schema)
+ #expect(try await df.schema == schema)
+ #expect(try await df.collect() == [Row(1, "Alice"), Row(2, nil)])
+
+ // Same result as the DDL string version.
+ let df2 = try await spark.createDataFrame(
+ [[1, "Alice"], [2, nil]], "id INT NOT NULL, name STRING")
+ #expect(try await df2.schema == schema)
+ #expect(try await df2.collect() == [Row(1, "Alice"), Row(2, nil)])
+ await spark.stop()
+ }
+
@Test
func supportedTypes() async throws {
let spark = try await SparkSession.builder.getOrCreate()
diff --git a/Tests/SparkConnectTests/DataFrameReaderTests.swift
b/Tests/SparkConnectTests/DataFrameReaderTests.swift
index a4ce325..71210c6 100644
--- a/Tests/SparkConnectTests/DataFrameReaderTests.swift
+++ b/Tests/SparkConnectTests/DataFrameReaderTests.swift
@@ -151,6 +151,19 @@ struct DataFrameReaderTests {
await spark.stop()
}
+ @Test
+ func schemaWithStructType() async throws {
+ let spark = try await SparkSession.builder.getOrCreate()
+ let path = "../examples/src/main/resources/people.json"
+ let structType = StructType(fields: [
+ StructField(name: "age", dataType: .short),
+ StructField(name: "name", dataType: .string),
+ ])
+ let expected = try await spark.read.schema("age SHORT, name
STRING").json(path).schema
+ #expect(try await spark.read.schema(structType).json(path).schema ==
expected)
+ await spark.stop()
+ }
+
@Test
func invalidSchema() async throws {
let spark = try await SparkSession.builder.getOrCreate()
diff --git a/Tests/SparkConnectTests/DataFrameTests.swift
b/Tests/SparkConnectTests/DataFrameTests.swift
index b67adbc..98ef0ee 100644
--- a/Tests/SparkConnectTests/DataFrameTests.swift
+++ b/Tests/SparkConnectTests/DataFrameTests.swift
@@ -316,6 +316,31 @@ struct DataFrameTests {
await spark.stop()
}
+ @Test
+ func toWithStructType() async throws {
+ let spark = try await SparkSession.builder.getOrCreate()
+ let df = try await spark.sql(
+ "SELECT 1 AS id, struct('a' AS x) AS s, array(1, 2) AS a, map('k', 1.0)
AS m")
+ let structType = StructType(fields: [
+ StructField(name: "id", dataType: .short, nullable: false),
+ StructField(
+ name: "s", dataType: .struct(StructType(fields: [StructField(name:
"x", dataType: .string)]))),
+ StructField(name: "a", dataType: .array(elementType: .long,
containsNull: true)),
+ StructField(
+ name: "m",
+ dataType: .map(
+ keyType: .string, valueType: .decimal(precision: 10, scale: 2),
valueContainsNull: true)),
+ ])
+ let ddl = "id SMALLINT NOT NULL,s STRUCT<x: STRING>,a ARRAY<BIGINT>,m
MAP<STRING, DECIMAL(10,2)>"
+ #expect(structType.toDDL == ddl)
+
+ let schema1 = try await df.to(ddl).schema
+ let schema2 = try await df.to(structType).schema
+ #expect(schema1 == schema2)
+
+ await spark.stop()
+ }
+
@Test
func selectMultipleColumns() async throws {
let spark = try await SparkSession.builder.getOrCreate()
diff --git a/Tests/SparkConnectTests/DataStreamTests.swift
b/Tests/SparkConnectTests/DataStreamTests.swift
index 038dfdf..4ec98b2 100644
--- a/Tests/SparkConnectTests/DataStreamTests.swift
+++ b/Tests/SparkConnectTests/DataStreamTests.swift
@@ -70,4 +70,18 @@ struct DataStreamTests {
#expect(try await df3.count() == 2025)
await spark.stop()
}
+
+ @Test
+ func schemaWithStructType() async throws {
+ let spark = try await SparkSession.builder.getOrCreate()
+ let input = "/tmp/input-" + UUID().uuidString
+ try await spark.range(1).write.orc(input)
+
+ let structType = StructType(fields: [StructField(name: "id", dataType:
.long)])
+ let df1 = try await spark.readStream.schema("id LONG").orc(input)
+ let df2 = try await spark.readStream.schema(structType).orc(input)
+ #expect(try await df2.isStreaming())
+ #expect(try await df1.schema == df2.schema)
+ await spark.stop()
+ }
}
diff --git a/Tests/SparkConnectTests/StructTypeTests.swift
b/Tests/SparkConnectTests/StructTypeTests.swift
index db20b4d..befb419 100644
--- a/Tests/SparkConnectTests/StructTypeTests.swift
+++ b/Tests/SparkConnectTests/StructTypeTests.swift
@@ -56,6 +56,31 @@ struct StructTypeTests {
#expect(count == 3)
}
+ @Test
+ func toDDL() {
+ let nested = StructType(fields: [
+ StructField(name: "x", dataType: .integer, nullable: false),
+ StructField(name: "y z", dataType: .string),
+ ])
+ let structType = StructType(fields: [
+ StructField(name: "id", dataType: .long, nullable: false),
+ StructField(name: "1st", dataType: .decimal(precision: 10, scale: 2)),
+ StructField(name: "t", dataType: .time(precision: 6)),
+ StructField(name: "s", dataType: .struct(nested)),
+ StructField(name: "arr", dataType: .array(elementType: .struct(nested),
containsNull: true)),
+ StructField(
+ name: "m", dataType: .map(keyType: .string, valueType: .date,
valueContainsNull: true)),
+ ])
+ #expect(
+ structType.toDDL
+ == "id BIGINT NOT NULL,`1st` DECIMAL(10,2),t TIME(6),"
+ + "s STRUCT<x: INT NOT NULL, `y z`: STRING>,"
+ + "arr ARRAY<STRUCT<x: INT NOT NULL, `y z`: STRING>>,m MAP<STRING,
DATE>")
+ #expect(StructType().toDDL == "")
+ #expect(StructField(name: "a`b", dataType: .string).toDDL == "`a``b`
STRING")
+ #expect(StructField(name: "", dataType: .string).toDDL == "`` STRING")
+ }
+
@Test
func equatable() {
let a = StructType(fields: [StructField(name: "id", dataType: .long)])
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
