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/08/07 18:04:06 UTC
(spark-connect-swift) branch main updated: [SPARK-58647] Add public `DataType` and `StructType`
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 b0e60bd [SPARK-58647] Add public `DataType` and `StructType`
b0e60bd is described below
commit b0e60bd28ab962b36859d48354a8c51df507f1b9
Author: Dongjoon Hyun <do...@apache.org>
AuthorDate: Fri Aug 7 11:03:59 2026 -0700
[SPARK-58647] Add public `DataType` and `StructType`
### What changes were proposed in this pull request?
This PR adds public schema model types — `DataType`, `StructType`, and `StructField` — and changes `DataFrame.schema` to return a `StructType` instead of a JSON string.
- `DataType` is an `indirect enum` covering all Spark Connect type kinds, with `simpleString` for Spark SQL-style representations like `decimal(10,2)` and `array<string>`.
- `StructType` holds an ordered list of `StructField`s (`name`/`dataType`/`nullable`/`metadata`) and supports subscripts and iteration. All types are `Sendable` and `Equatable`.
```swift
let schema = try await df.schema
for field in schema {
print("\(field.name): \(field.dataType.simpleString) (nullable=\(field.nullable))")
}
```
### Why are the changes needed?
For feature parity with Scala/PySpark where `df.schema` returns a traversable `StructType` instead of a raw JSON string.
### Does this PR introduce _any_ user-facing change?
Yes. `DataFrame.schema` now returns a `StructType` instead of a JSON `String`.
### How was this patch tested?
Pass the CIs with newly added `DataTypeTests` and `StructTypeTests` suites and the revised `DataFrameTests`.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Fable 5
Closes #487 from dongjoon-hyun/SPARK-58647.
Authored-by: Dongjoon Hyun <do...@apache.org>
Signed-off-by: Dongjoon Hyun <do...@apache.org>
---
Sources/SparkConnect/ConvertToArrow.swift | 6 +-
Sources/SparkConnect/DataFrame+Actions.swift | 2 +-
Sources/SparkConnect/DataFrame.swift | 27 +--
Sources/SparkConnect/DataType.swift | 267 +++++++++++++++++++++
.../Documentation.docc/SparkConnect.md | 9 +
Sources/SparkConnect/Extension.swift | 142 +----------
Sources/SparkConnect/StructType.swift | 82 +++++++
Sources/SparkConnect/TypeAliases.swift | 7 +-
Tests/SparkConnectTests/DataFrameTests.swift | 55 ++---
Tests/SparkConnectTests/DataTypeTests.swift | 184 ++++++++++++++
.../SparkConnectClientTests.swift | 2 +-
Tests/SparkConnectTests/StructTypeTests.swift | 71 ++++++
12 files changed, 659 insertions(+), 195 deletions(-)
diff --git a/Sources/SparkConnect/ConvertToArrow.swift b/Sources/SparkConnect/ConvertToArrow.swift
index c626e7e..0bc0515 100644
--- a/Sources/SparkConnect/ConvertToArrow.swift
+++ b/Sources/SparkConnect/ConvertToArrow.swift
@@ -28,9 +28,9 @@ enum ConvertToArrow {
/// Convert the given rows into an `Apache Arrow` IPC stream according to the Spark schema.
/// - Parameters:
/// - data: An array of rows whose values are ordered like the schema fields.
- /// - schema: A ``StructType`` describing the field names and types.
+ /// - schema: A `ProtoStructType` describing the field names and types.
/// - Returns: A serialized `Apache Arrow` IPC stream.
- static func toArrowIPCStream(_ data: [[Sendable?]], _ schema: StructType) throws -> Data {
+ static func toArrowIPCStream(_ data: [[Sendable?]], _ schema: ProtoStructType) throws -> Data {
for row in data where row.count != schema.fields.count {
throw SparkConnectError.InvalidArgument
}
@@ -57,7 +57,7 @@ enum ConvertToArrow {
}
/// Build an Arrow column from the given values according to the Spark data type.
- private static func toArrowColumn(_ column: [Sendable?], _ dataType: DataType) throws
+ private static func toArrowColumn(_ column: [Sendable?], _ dataType: ProtoDataType) throws
-> ArrowArrayHolder
{
switch dataType.kind {
diff --git a/Sources/SparkConnect/DataFrame+Actions.swift b/Sources/SparkConnect/DataFrame+Actions.swift
index 3b9c20f..517438d 100644
--- a/Sources/SparkConnect/DataFrame+Actions.swift
+++ b/Sources/SparkConnect/DataFrame+Actions.swift
@@ -43,7 +43,7 @@ extension DataFrame {
m in
if m.hasSchema {
// The original schema should arrive before ArrowBatches
- await self.setSchema(m.schema)
+ try await self.setSchema(DataType(m.schema))
}
let ipcStreamBytes = m.arrowBatch.data
if !ipcStreamBytes.isEmpty && m.arrowBatch.rowCount > 0 {
diff --git a/Sources/SparkConnect/DataFrame.swift b/Sources/SparkConnect/DataFrame.swift
index 1dfa024..03a20ce 100644
--- a/Sources/SparkConnect/DataFrame.swift
+++ b/Sources/SparkConnect/DataFrame.swift
@@ -202,7 +202,7 @@ import Synchronization
public actor DataFrame: Sendable {
var spark: SparkSession
var plan: Plan
- var _schema: DataType? = nil
+ var _schema: StructType? = nil
var batches: [RecordBatch] = [RecordBatch]()
/// Create a new `DataFrame`instance with the given Spark session and plan.
@@ -249,9 +249,12 @@ public actor DataFrame: Sendable {
}
/// Set the schema. This is used to store the analized schema response from `Spark Connect` server.
- /// - Parameter schema: <#schema description#>
- func setSchema(_ schema: DataType) {
- self._schema = schema
+ /// - Parameter schema: A ``DataType`` which is expected to be a struct.
+ func setSchema(_ schema: DataType) throws {
+ guard case .struct(let structType) = schema else {
+ throw SparkConnectError.InvalidType
+ }
+ self._schema = structType
}
/// Add `Apache Arrow`'s `RecordBatch`s to the internal array.
@@ -279,20 +282,16 @@ public actor DataFrame: Sendable {
/// Return an array of column name strings
public var columns: [String] {
get async throws {
- var columns: [String] = []
try await analyzePlanIfNeeded()
- for field in self._schema!.struct.fields {
- columns.append(field.name)
- }
- return columns
+ return self._schema!.fieldNames
}
}
- /// Return a `JSON` string of data type because we cannot expose the internal type ``DataType``.
- public var schema: String {
+ /// Return the schema of this ``DataFrame`` as a ``StructType``.
+ public var schema: StructType {
get async throws {
try await analyzePlanIfNeeded()
- return try self._schema!.jsonString()
+ return self._schema!
}
}
@@ -300,7 +299,7 @@ public actor DataFrame: Sendable {
public var dtypes: [(String, String)] {
get async throws {
try await analyzePlanIfNeeded()
- return try self._schema!.struct.fields.map { ($0.name, try $0.dataType.simpleString) }
+ return self._schema!.fields.map { ($0.name, $0.dataType.simpleString) }
}
}
@@ -334,7 +333,7 @@ public actor DataFrame: Sendable {
let service = Spark_Connect_SparkConnectService.Client(wrapping: client)
let response = try await service.analyzePlan(
spark.client.getAnalyzePlanRequest(spark.sessionID, plan))
- self.setSchema(response.schema.schema)
+ try self.setSchema(DataType(response.schema.schema))
}
}
diff --git a/Sources/SparkConnect/DataType.swift b/Sources/SparkConnect/DataType.swift
new file mode 100644
index 0000000..7fe7f54
--- /dev/null
+++ b/Sources/SparkConnect/DataType.swift
@@ -0,0 +1,267 @@
+//
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you 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 field of Spark's `YEAR-MONTH INTERVAL` type.
+public enum YearMonthIntervalField: Int32, Sendable, Equatable {
+ case year = 0
+ case month = 1
+
+ var fieldName: String {
+ switch self {
+ case .year: "year"
+ case .month: "month"
+ }
+ }
+}
+
+/// A field of Spark's `DAY-TIME INTERVAL` type.
+public enum DayTimeIntervalField: Int32, Sendable, Equatable {
+ case day = 0
+ case hour = 1
+ case minute = 2
+ case second = 3
+
+ var fieldName: String {
+ switch self {
+ case .day: "day"
+ case .hour: "hour"
+ case .minute: "minute"
+ case .second: "second"
+ }
+ }
+}
+
+/// A user-defined type (UDT) definition delivered by the `Spark Connect` server.
+public struct UserDefinedType: Sendable, Equatable {
+ /// The JVM class implementing this UDT, if any.
+ public var jvmClass: String?
+ /// The Python class implementing this UDT, if any.
+ public var pythonClass: String?
+ /// The serialized Python class of this UDT, if any.
+ public var serializedPythonClass: String?
+ /// The underlying SQL type of this UDT, if any.
+ public var sqlType: DataType?
+
+ public init(
+ jvmClass: String? = nil, pythonClass: String? = nil, serializedPythonClass: String? = nil,
+ sqlType: DataType? = nil
+ ) {
+ self.jvmClass = jvmClass
+ self.pythonClass = pythonClass
+ self.serializedPythonClass = serializedPythonClass
+ self.sqlType = sqlType
+ }
+}
+
+/// The data type of a ``DataFrame`` column or of a ``StructField``, mirroring Spark SQL's
+/// `org.apache.spark.sql.types.DataType` hierarchy as a Swift enum.
+public indirect enum DataType: Sendable, Equatable {
+ case null
+ case binary
+ case boolean
+ case byte
+ case short
+ case integer
+ case long
+ case float
+ case double
+ case decimal(precision: Int32, scale: Int32)
+ case string
+ case char(length: Int32)
+ case varchar(length: Int32)
+ case date
+ case timestamp
+ case timestampNtz
+ case time(precision: Int32)
+ case calendarInterval
+ case yearMonthInterval(startField: YearMonthIntervalField, endField: YearMonthIntervalField)
+ case dayTimeInterval(startField: DayTimeIntervalField, endField: DayTimeIntervalField)
+ case array(elementType: DataType, containsNull: Bool)
+ case map(keyType: DataType, valueType: DataType, valueContainsNull: Bool)
+ case `struct`(StructType)
+ case variant
+ case geometry(srid: Int32)
+ case geography(srid: Int32)
+ case udt(UserDefinedType)
+ case unparsed(String)
+
+ /// A compact string representation like Spark SQL's `DataType.simpleString`,
+ /// e.g. `int`, `decimal(10,2)`, or `array<string>`.
+ public var simpleString: String {
+ switch self {
+ case .null:
+ "void"
+ case .binary:
+ "binary"
+ case .boolean:
+ "boolean"
+ case .byte:
+ "tinyint"
+ case .short:
+ "smallint"
+ case .integer:
+ "int"
+ case .long:
+ "bigint"
+ case .float:
+ "float"
+ case .double:
+ "double"
+ case .decimal(let precision, let scale):
+ "decimal(\(precision),\(scale))"
+ case .string:
+ "string"
+ case .char(let length):
+ "char(\(length))"
+ case .varchar(let length):
+ "varchar(\(length))"
+ case .date:
+ "date"
+ case .timestamp:
+ "timestamp"
+ case .timestampNtz:
+ "timestamp_ntz"
+ case .time(let precision):
+ "time(\(precision))"
+ case .calendarInterval:
+ "interval"
+ case .yearMonthInterval(let startField, let endField):
+ if startField == endField {
+ "interval \(startField.fieldName)"
+ } else {
+ "interval \(startField.fieldName) to \(endField.fieldName)"
+ }
+ case .dayTimeInterval(let startField, let endField):
+ if startField == endField {
+ "interval \(startField.fieldName)"
+ } else {
+ "interval \(startField.fieldName) to \(endField.fieldName)"
+ }
+ case .array(let elementType, _):
+ "array<\(elementType.simpleString)>"
+ case .map(let keyType, let valueType, _):
+ "map<\(keyType.simpleString),\(valueType.simpleString)>"
+ case .struct(let structType):
+ structType.simpleString
+ case .variant:
+ "variant"
+ case .geometry(let srid):
+ if srid == -1 {
+ "geometry(any)"
+ } else {
+ "geometry(\(srid))"
+ }
+ case .geography(let srid):
+ if srid == -1 {
+ "geography(any)"
+ } else {
+ "geography(\(srid))"
+ }
+ case .udt:
+ "udt"
+ case .unparsed(let dataTypeString):
+ dataTypeString
+ }
+ }
+}
+
+extension DataType {
+ /// Create a public ``DataType`` from the `Spark Connect` protobuf representation.
+ init(_ proto: ProtoDataType) throws {
+ switch proto.kind {
+ case .null:
+ self = .null
+ case .binary:
+ self = .binary
+ case .boolean:
+ self = .boolean
+ case .byte:
+ self = .byte
+ case .short:
+ self = .short
+ case .integer:
+ self = .integer
+ case .long:
+ self = .long
+ case .float:
+ self = .float
+ case .double:
+ self = .double
+ case .decimal(let decimal):
+ self = .decimal(
+ precision: decimal.hasPrecision ? decimal.precision : 10,
+ scale: decimal.hasScale ? decimal.scale : 0)
+ case .string:
+ self = .string
+ case .char(let char):
+ self = .char(length: char.length)
+ case .varChar(let varChar):
+ self = .varchar(length: varChar.length)
+ case .date:
+ self = .date
+ case .timestamp:
+ self = .timestamp
+ case .timestampNtz:
+ self = .timestampNtz
+ case .time(let time):
+ self = .time(precision: time.hasPrecision ? time.precision : 6)
+ case .calendarInterval:
+ self = .calendarInterval
+ case .yearMonthInterval(let interval):
+ guard let startField = YearMonthIntervalField(rawValue: interval.startField),
+ let endField = YearMonthIntervalField(rawValue: interval.endField)
+ else {
+ throw SparkConnectError.InvalidType
+ }
+ self = .yearMonthInterval(startField: startField, endField: endField)
+ case .dayTimeInterval(let interval):
+ guard let startField = DayTimeIntervalField(rawValue: interval.startField),
+ let endField = DayTimeIntervalField(rawValue: interval.endField)
+ else {
+ throw SparkConnectError.InvalidType
+ }
+ self = .dayTimeInterval(startField: startField, endField: endField)
+ case .array(let array):
+ self = .array(
+ elementType: try DataType(array.elementType), containsNull: array.containsNull)
+ case .map(let map):
+ self = .map(
+ keyType: try DataType(map.keyType), valueType: try DataType(map.valueType),
+ valueContainsNull: map.valueContainsNull)
+ case .struct(let structType):
+ self = .struct(try StructType(structType))
+ case .variant:
+ self = .variant
+ case .geometry(let geometry):
+ self = .geometry(srid: geometry.srid)
+ case .geography(let geography):
+ self = .geography(srid: geography.srid)
+ case .udt(let udt):
+ self = .udt(
+ UserDefinedType(
+ jvmClass: udt.hasJvmClass ? udt.jvmClass : nil,
+ pythonClass: udt.hasPythonClass ? udt.pythonClass : nil,
+ serializedPythonClass: udt.hasSerializedPythonClass ? udt.serializedPythonClass : nil,
+ sqlType: udt.hasSqlType ? try DataType(udt.sqlType) : nil))
+ case .unparsed(let unparsed):
+ self = .unparsed(unparsed.dataTypeString)
+ case .none:
+ throw SparkConnectError.InvalidType
+ }
+ }
+}
diff --git a/Sources/SparkConnect/Documentation.docc/SparkConnect.md b/Sources/SparkConnect/Documentation.docc/SparkConnect.md
index 3d6bc5a..cfd40fb 100644
--- a/Sources/SparkConnect/Documentation.docc/SparkConnect.md
+++ b/Sources/SparkConnect/Documentation.docc/SparkConnect.md
@@ -33,6 +33,15 @@ SparkConnect is a modern Swift library that provides a native interface to Apach
- ``LocalTime``
- ``StorageLevel``
+### Data Types
+
+- ``DataType``
+- ``StructType``
+- ``StructField``
+- ``UserDefinedType``
+- ``YearMonthIntervalField``
+- ``DayTimeIntervalField``
+
### Data I/O
- ``DataFrameReader``
diff --git a/Sources/SparkConnect/Extension.swift b/Sources/SparkConnect/Extension.swift
index 7983ba7..01354c4 100644
--- a/Sources/SparkConnect/Extension.swift
+++ b/Sources/SparkConnect/Extension.swift
@@ -202,8 +202,8 @@ extension ExpressionLiteral {
self.time.precision = 6
default:
if case Optional<Any>.none = value as Any {
- var dataType = DataType()
- dataType.null = DataType.NULL()
+ var dataType = ProtoDataType()
+ dataType.null = ProtoDataType.NULL()
self.null = dataType
} else {
throw SparkConnectError.InvalidType
@@ -259,141 +259,3 @@ extension SparkSession: Equatable {
return lhs.sessionID == rhs.sessionID
}
}
-
-protocol IntervalType {
- var startField: Int32 { get }
- var endField: Int32 { get }
- func fieldToString(_ field: Int32) throws -> String
-}
-
-extension IntervalType {
- func toString() throws -> String {
- let startFieldName = try fieldToString(self.startField)
- let endFieldName = try fieldToString(self.endField)
- let interval =
- if startFieldName == endFieldName {
- "interval \(startFieldName)"
- } else if startFieldName < endFieldName {
- "interval \(startFieldName) to \(endFieldName)"
- } else {
- throw SparkConnectError.InvalidType
- }
- return interval
- }
-}
-
-extension YearMonthInterval: IntervalType {
- func fieldToString(_ field: Int32) throws -> String {
- return switch field {
- case 0: "year"
- case 1: "month"
- default:
- throw SparkConnectError.InvalidType
- }
- }
-}
-
-extension DayTimeInterval: IntervalType {
- func fieldToString(_ field: Int32) throws -> String {
- return switch field {
- case 0: "day"
- case 1: "hour"
- case 2: "minute"
- case 3: "second"
- default:
- throw SparkConnectError.InvalidType
- }
- }
-}
-
-extension MapType {
- func toString() throws -> String {
- return "map<\(try self.keyType.simpleString),\(try self.valueType.simpleString)>"
- }
-}
-
-extension StructType {
- func toString() throws -> String {
- let fieldTypes = try fields.map { "\($0.name):\(try $0.dataType.simpleString)" }
- return "struct<\(fieldTypes.joined(separator: ","))>"
- }
-}
-
-extension DataType {
- var simpleString: String {
- get throws {
- return switch self.kind {
- case .null:
- "void"
- case .binary:
- "binary"
- case .boolean:
- "boolean"
- case .byte:
- "tinyint"
- case .short:
- "smallint"
- case .integer:
- "int"
- case .long:
- "bigint"
- case .float:
- "float"
- case .double:
- "double"
- case .decimal:
- "decimal(\(self.decimal.precision),\(self.decimal.scale))"
- case .string:
- "string"
- case .char:
- "char"
- case .varChar:
- "varchar"
- case .date:
- "date"
- case .timestamp:
- "timestamp"
- case .timestampNtz:
- "timestamp_ntz"
- case .time:
- if self.time.hasPrecision {
- "time(\(self.time.precision))"
- } else {
- "time(6)"
- }
- case .calendarInterval:
- "interval"
- case .yearMonthInterval:
- try self.yearMonthInterval.toString()
- case .dayTimeInterval:
- try self.dayTimeInterval.toString()
- case .array:
- "array<\(try self.array.elementType.simpleString)>"
- case .struct:
- try self.struct.toString()
- case .map:
- try self.map.toString()
- case .variant:
- "variant"
- case .geometry:
- if self.geometry.srid == -1 {
- "geometry(any)"
- } else {
- "geometry(\(self.geometry.srid))"
- }
- case .geography:
- if self.geography.srid == -1 {
- "geography(any)"
- } else {
- "geography(\(self.geography.srid))"
- }
- case .udt:
- self.udt.type
- case .unparsed:
- self.unparsed.dataTypeString
- default:
- throw SparkConnectError.InvalidType
- }
- }
- }
-}
diff --git a/Sources/SparkConnect/StructType.swift b/Sources/SparkConnect/StructType.swift
new file mode 100644
index 0000000..3896230
--- /dev/null
+++ b/Sources/SparkConnect/StructType.swift
@@ -0,0 +1,82 @@
+//
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you 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 field inside a ``StructType`` with a name, a ``DataType``, a nullability flag, and
+/// an optional JSON-encoded metadata string.
+public struct StructField: Sendable, Equatable {
+ public var name: String
+ public var dataType: DataType
+ public var nullable: Bool
+ public var metadata: String?
+
+ public init(name: String, dataType: DataType, nullable: Bool = true, metadata: String? = nil) {
+ self.name = name
+ self.dataType = dataType
+ self.nullable = nullable
+ self.metadata = metadata
+ }
+}
+
+/// A struct type holding an ordered collection of ``StructField``s, mirroring Spark SQL's
+/// `org.apache.spark.sql.types.StructType`. This is the schema representation of
+/// a ``DataFrame`` returned by ``DataFrame/schema``.
+public struct StructType: Sendable, Equatable {
+ public var fields: [StructField]
+
+ public init(fields: [StructField] = []) {
+ self.fields = fields
+ }
+
+ /// The names of all fields in order.
+ public var fieldNames: [String] {
+ fields.map { $0.name }
+ }
+
+ /// The first field whose name is equal to the given name, or nil if it doesn't exist.
+ public subscript(name: String) -> StructField? {
+ fields.first { $0.name == name }
+ }
+
+ /// A compact string representation like Spark SQL's `StructType.simpleString`,
+ /// e.g. `struct<id:bigint,name:string>`.
+ public var simpleString: String {
+ "struct<\(fields.map { "\($0.name):\($0.dataType.simpleString)" }.joined(separator: ","))>"
+ }
+}
+
+extension StructType: RandomAccessCollection {
+ public var startIndex: Int { fields.startIndex }
+ public var endIndex: Int { fields.endIndex }
+
+ public subscript(position: Int) -> StructField {
+ fields[position]
+ }
+}
+
+extension StructType {
+ /// Create a public ``StructType`` from the `Spark Connect` protobuf representation.
+ init(_ proto: ProtoStructType) throws {
+ self.fields = try proto.fields.map {
+ StructField(
+ name: $0.name,
+ dataType: try DataType($0.dataType),
+ nullable: $0.nullable,
+ metadata: $0.hasMetadata ? $0.metadata : nil)
+ }
+ }
+}
diff --git a/Sources/SparkConnect/TypeAliases.swift b/Sources/SparkConnect/TypeAliases.swift
index f7f4d83..43c8a00 100644
--- a/Sources/SparkConnect/TypeAliases.swift
+++ b/Sources/SparkConnect/TypeAliases.swift
@@ -26,8 +26,6 @@ typealias CollectMetrics = Spark_Connect_CollectMetrics
typealias Command = Spark_Connect_Command
typealias ConfigRequest = Spark_Connect_ConfigRequest
typealias DataSource = Spark_Connect_Read.DataSource
-typealias DataType = Spark_Connect_DataType
-typealias DayTimeInterval = Spark_Connect_DataType.DayTimeInterval
typealias Drop = Spark_Connect_Drop
typealias ExecutePlanRequest = Spark_Connect_ExecutePlanRequest
typealias ExecutePlanResponse = Spark_Connect_ExecutePlanResponse
@@ -43,7 +41,6 @@ typealias JoinType = Spark_Connect_Join.JoinType
typealias KeyValue = Spark_Connect_KeyValue
typealias LateralJoin = Spark_Connect_LateralJoin
typealias Limit = Spark_Connect_Limit
-typealias MapType = Spark_Connect_DataType.Map
typealias MergeAction = Spark_Connect_MergeAction
typealias MergeIntoTableCommand = Spark_Connect_MergeIntoTableCommand
typealias NamedTable = Spark_Connect_Read.NamedTable
@@ -54,6 +51,8 @@ typealias OutputType = Spark_Connect_OutputType
typealias Parse = Spark_Connect_Parse
typealias Plan = Spark_Connect_Plan
typealias Project = Spark_Connect_Project
+typealias ProtoDataType = Spark_Connect_DataType
+typealias ProtoStructType = Spark_Connect_DataType.Struct
typealias Range = Spark_Connect_Range
typealias Read = Spark_Connect_Read
typealias ReattachExecuteRequest = Spark_Connect_ReattachExecuteRequest
@@ -72,7 +71,6 @@ typealias Sort = Spark_Connect_Sort
typealias StreamingQueryCommand = Spark_Connect_StreamingQueryCommand
typealias StreamingQueryInstanceId = Spark_Connect_StreamingQueryInstanceId
typealias StreamingQueryManagerCommand = Spark_Connect_StreamingQueryManagerCommand
-typealias StructType = Spark_Connect_DataType.Struct
typealias Tail = Spark_Connect_Tail
typealias UserContext = Spark_Connect_UserContext
typealias UnresolvedAttribute = Spark_Connect_Expression.UnresolvedAttribute
@@ -81,4 +79,3 @@ typealias WithColumnsRenamed = Spark_Connect_WithColumnsRenamed
typealias WriteOperation = Spark_Connect_WriteOperation
typealias WriteOperationV2 = Spark_Connect_WriteOperationV2
typealias WriteStreamOperationStart = Spark_Connect_WriteStreamOperationStart
-typealias YearMonthInterval = Spark_Connect_DataType.YearMonthInterval
diff --git a/Tests/SparkConnectTests/DataFrameTests.swift b/Tests/SparkConnectTests/DataFrameTests.swift
index 261b58b..b67adbc 100644
--- a/Tests/SparkConnectTests/DataFrameTests.swift
+++ b/Tests/SparkConnectTests/DataFrameTests.swift
@@ -73,26 +73,25 @@ struct DataFrameTests {
let spark = try await SparkSession.builder.getOrCreate()
let schema1 = try await spark.sql("SELECT 'a' as col1").schema
- let version = await spark.version
- let answer1 =
- if ["4.0", "4.1"].contains(where: version.hasPrefix) {
- #"{"struct":{"fields":[{"name":"col1","dataType":{"string":{"collation":"UTF8_BINARY"}}}]}}"#
- } else {
- #"{"struct":{"fields":[{"name":"col1","dataType":{"string":{}}}]}}"#
- }
- #expect(schema1 == answer1)
+ #expect(
+ schema1
+ == StructType(fields: [StructField(name: "col1", dataType: .string, nullable: false)]))
+ #expect(schema1.fieldNames == ["col1"])
+ #expect(schema1["col1"]?.dataType == .string)
+ #expect(schema1.count == 1)
let schema2 = try await spark.sql("SELECT 'a' as col1, 'b' as col2").schema
- let answer2 =
- if ["4.0", "4.1"].contains(where: version.hasPrefix) {
- #"{"struct":{"fields":[{"name":"col1","dataType":{"string":{"collation":"UTF8_BINARY"}}},{"name":"col2","dataType":{"string":{"collation":"UTF8_BINARY"}}}]}}"#
- } else {
- #"{"struct":{"fields":[{"name":"col1","dataType":{"string":{}}},{"name":"col2","dataType":{"string":{}}}]}}"#
- }
- #expect(schema2 == answer2)
+ #expect(
+ schema2
+ == StructType(fields: [
+ StructField(name: "col1", dataType: .string, nullable: false),
+ StructField(name: "col2", dataType: .string, nullable: false),
+ ]))
+ #expect(schema2.map { $0.name } == ["col1", "col2"])
let emptySchema = try await spark.sql("DROP TABLE IF EXISTS nonexistent").schema
- #expect(emptySchema == #"{"struct":{}}"#)
+ #expect(emptySchema == StructType())
+ #expect(emptySchema.isEmpty)
await spark.stop()
}
@@ -279,7 +278,7 @@ struct DataFrameTests {
func selectNone() async throws {
let spark = try await SparkSession.builder.getOrCreate()
let emptySchema = try await spark.range(1).select().schema
- #expect(emptySchema == #"{"struct":{}}"#)
+ #expect(emptySchema == StructType())
await spark.stop()
}
@@ -289,9 +288,7 @@ struct DataFrameTests {
#expect(try await spark.range(1).select().columns.isEmpty)
let schema = try await spark.range(1).select("id").schema
#expect(
- schema
- == #"{"struct":{"fields":[{"name":"id","dataType":{"long":{}}}]}}"#
- )
+ schema == StructType(fields: [StructField(name: "id", dataType: .long, nullable: false)]))
await spark.stop()
}
@@ -309,16 +306,12 @@ struct DataFrameTests {
let schema1 = try await spark.range(1).to("shortID SHORT").schema
#expect(
- schema1
- == #"{"struct":{"fields":[{"name":"shortID","dataType":{"short":{}},"nullable":true}]}}"#
- )
+ schema1 == StructType(fields: [StructField(name: "shortID", dataType: .short)]))
let schema2 = try await spark.sql("SELECT '1'").to("id INT").schema
print(schema2)
#expect(
- schema2
- == #"{"struct":{"fields":[{"name":"id","dataType":{"integer":{}},"nullable":true}]}}"#
- )
+ schema2 == StructType(fields: [StructField(name: "id", dataType: .integer)]))
await spark.stop()
}
@@ -329,8 +322,10 @@ struct DataFrameTests {
let schema = try await spark.sql("SELECT * FROM VALUES (1, 2)").select("col2", "col1").schema
#expect(
schema
- == #"{"struct":{"fields":[{"name":"col2","dataType":{"integer":{}}},{"name":"col1","dataType":{"integer":{}}}]}}"#
- )
+ == StructType(fields: [
+ StructField(name: "col2", dataType: .integer, nullable: false),
+ StructField(name: "col1", dataType: .integer, nullable: false),
+ ]))
await spark.stop()
}
@@ -351,9 +346,7 @@ struct DataFrameTests {
let spark = try await SparkSession.builder.getOrCreate()
let schema = try await spark.range(1).selectExpr("id + 1 as id2").schema
#expect(
- schema
- == #"{"struct":{"fields":[{"name":"id2","dataType":{"long":{}}}]}}"#
- )
+ schema == StructType(fields: [StructField(name: "id2", dataType: .long, nullable: false)]))
await spark.stop()
}
diff --git a/Tests/SparkConnectTests/DataTypeTests.swift b/Tests/SparkConnectTests/DataTypeTests.swift
new file mode 100644
index 0000000..a732e69
--- /dev/null
+++ b/Tests/SparkConnectTests/DataTypeTests.swift
@@ -0,0 +1,184 @@
+//
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you 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.
+//
+
+import SparkConnect
+import Testing
+
+/// A test suite for `DataType`, `StructField`, and `StructType`
+@Suite(.serialized)
+struct DataTypeTests {
+ @Test
+ func simpleString() {
+ #expect(DataType.null.simpleString == "void")
+ #expect(DataType.binary.simpleString == "binary")
+ #expect(DataType.boolean.simpleString == "boolean")
+ #expect(DataType.byte.simpleString == "tinyint")
+ #expect(DataType.short.simpleString == "smallint")
+ #expect(DataType.integer.simpleString == "int")
+ #expect(DataType.long.simpleString == "bigint")
+ #expect(DataType.float.simpleString == "float")
+ #expect(DataType.double.simpleString == "double")
+ #expect(DataType.decimal(precision: 10, scale: 2).simpleString == "decimal(10,2)")
+ #expect(DataType.string.simpleString == "string")
+ #expect(DataType.char(length: 5).simpleString == "char(5)")
+ #expect(DataType.varchar(length: 10).simpleString == "varchar(10)")
+ #expect(DataType.date.simpleString == "date")
+ #expect(DataType.timestamp.simpleString == "timestamp")
+ #expect(DataType.timestampNtz.simpleString == "timestamp_ntz")
+ #expect(DataType.time(precision: 6).simpleString == "time(6)")
+ #expect(DataType.calendarInterval.simpleString == "interval")
+ #expect(DataType.variant.simpleString == "variant")
+ #expect(DataType.geometry(srid: -1).simpleString == "geometry(any)")
+ #expect(DataType.geometry(srid: 4326).simpleString == "geometry(4326)")
+ #expect(DataType.geography(srid: -1).simpleString == "geography(any)")
+ #expect(DataType.geography(srid: 4326).simpleString == "geography(4326)")
+ #expect(DataType.udt(UserDefinedType()).simpleString == "udt")
+ #expect(DataType.unparsed("unknown").simpleString == "unknown")
+ }
+
+ @Test
+ func intervalSimpleString() {
+ #expect(
+ DataType.yearMonthInterval(startField: .year, endField: .year).simpleString
+ == "interval year")
+ #expect(
+ DataType.yearMonthInterval(startField: .year, endField: .month).simpleString
+ == "interval year to month")
+ #expect(
+ DataType.dayTimeInterval(startField: .day, endField: .day).simpleString == "interval day")
+ #expect(
+ DataType.dayTimeInterval(startField: .day, endField: .second).simpleString
+ == "interval day to second")
+ #expect(
+ DataType.dayTimeInterval(startField: .hour, endField: .minute).simpleString
+ == "interval hour to minute")
+ }
+
+ @Test
+ func nestedSimpleString() {
+ #expect(
+ DataType.array(elementType: .integer, containsNull: true).simpleString == "array<int>")
+ #expect(
+ DataType.map(keyType: .string, valueType: .long, valueContainsNull: true).simpleString
+ == "map<string,bigint>")
+ let structType = StructType(fields: [
+ StructField(name: "id", dataType: .long),
+ StructField(name: "arr", dataType: .array(elementType: .string, containsNull: false)),
+ ])
+ #expect(DataType.struct(structType).simpleString == "struct<id:bigint,arr:array<string>>")
+ }
+
+ @Test
+ func schema() async throws {
+ let spark = try await SparkSession.builder.getOrCreate()
+ let df = try await spark.sql(
+ """
+ SELECT
+ CAST(1 AS TINYINT) t,
+ CAST(1 AS SMALLINT) s,
+ CAST(1 AS INT) i,
+ CAST(1 AS BIGINT) l,
+ CAST(1.0 AS FLOAT) f,
+ CAST(1.0 AS DOUBLE) d,
+ CAST(1.5 AS DECIMAL(10,2)) dec,
+ 'a' str,
+ CAST('a' AS BINARY) bin,
+ true b,
+ DATE'2025-01-01' dt,
+ TIMESTAMP'2025-01-01 00:00:00' ts
+ """)
+ let schema = try await df.schema
+ #expect(schema.count == 12)
+ #expect(schema["t"]?.dataType == .byte)
+ #expect(schema["s"]?.dataType == .short)
+ #expect(schema["i"]?.dataType == .integer)
+ #expect(schema["l"]?.dataType == .long)
+ #expect(schema["f"]?.dataType == .float)
+ #expect(schema["d"]?.dataType == .double)
+ #expect(schema["dec"]?.dataType == .decimal(precision: 10, scale: 2))
+ #expect(schema["str"]?.dataType == .string)
+ #expect(schema["bin"]?.dataType == .binary)
+ #expect(schema["b"]?.dataType == .boolean)
+ #expect(schema["dt"]?.dataType == .date)
+ #expect(schema["ts"]?.dataType == .timestamp)
+ await spark.stop()
+ }
+
+ @Test
+ func schemaComplexTypes() async throws {
+ let spark = try await SparkSession.builder.getOrCreate()
+ let df = try await spark.sql(
+ """
+ SELECT
+ ARRAY(1, 2) arr,
+ MAP('k', 1) m,
+ NAMED_STRUCT('a', 1, 'b', ARRAY('x')) st
+ """)
+ let schema = try await df.schema
+ #expect(schema.fieldNames == ["arr", "m", "st"])
+ #expect(schema["arr"]?.dataType == .array(elementType: .integer, containsNull: false))
+ #expect(
+ schema["m"]?.dataType
+ == .map(keyType: .string, valueType: .integer, valueContainsNull: false))
+ #expect(
+ schema["st"]?.dataType
+ == .struct(
+ StructType(fields: [
+ StructField(name: "a", dataType: .integer, nullable: false),
+ StructField(
+ name: "b", dataType: .array(elementType: .string, containsNull: false),
+ nullable: false),
+ ])))
+ await spark.stop()
+ }
+
+ @Test
+ func schemaIntervalTypes() async throws {
+ let spark = try await SparkSession.builder.getOrCreate()
+ let df = try await spark.sql(
+ "SELECT INTERVAL '1' YEAR ym, INTERVAL '1 2:3:4' DAY TO SECOND dts")
+ let schema = try await df.schema
+ #expect(schema["ym"]?.dataType == .yearMonthInterval(startField: .year, endField: .year))
+ #expect(schema["dts"]?.dataType == .dayTimeInterval(startField: .day, endField: .second))
+ await spark.stop()
+ }
+
+ @Test
+ func schemaTimeType() async throws {
+ let spark = try await SparkSession.builder.getOrCreate()
+ if await spark.version >= "4.2" {
+ try await spark.conf.set("spark.sql.timeType.enabled", "true")
+ let df = try await spark.sql(
+ "SELECT TIME'12:34:56' t6, CAST(TIME'12:34:56' AS TIME(0)) t0")
+ let schema = try await df.schema
+ #expect(schema["t6"]?.dataType == .time(precision: 6))
+ #expect(schema["t0"]?.dataType == .time(precision: 0))
+ }
+ await spark.stop()
+ }
+
+ @Test
+ func schemaNullType() async throws {
+ let spark = try await SparkSession.builder.getOrCreate()
+ let schema = try await spark.sql("SELECT NULL n").schema
+ #expect(schema["n"]?.dataType == .null)
+ #expect(schema["n"]?.nullable == true)
+ await spark.stop()
+ }
+}
diff --git a/Tests/SparkConnectTests/SparkConnectClientTests.swift b/Tests/SparkConnectTests/SparkConnectClientTests.swift
index 36fd5e9..3b11777 100644
--- a/Tests/SparkConnectTests/SparkConnectClientTests.swift
+++ b/Tests/SparkConnectTests/SparkConnectClientTests.swift
@@ -112,7 +112,7 @@ struct SparkConnectClientTests {
func ddlParse() async throws {
let client = try SparkConnectClient(remote: TEST_REMOTE)
try await client.connect(UUID().uuidString)
- #expect(try await client.ddlParse("a int").simpleString == "struct<a:int>")
+ #expect(try await DataType(client.ddlParse("a int")).simpleString == "struct<a:int>")
await client.stop()
}
diff --git a/Tests/SparkConnectTests/StructTypeTests.swift b/Tests/SparkConnectTests/StructTypeTests.swift
new file mode 100644
index 0000000..db20b4d
--- /dev/null
+++ b/Tests/SparkConnectTests/StructTypeTests.swift
@@ -0,0 +1,71 @@
+//
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you 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.
+//
+
+import SparkConnect
+import Testing
+
+/// A test suite for `StructType` and `StructField`
+@Suite(.serialized)
+struct StructTypeTests {
+ @Test
+ func simpleString() {
+ let structType = StructType(fields: [
+ StructField(name: "id", dataType: .long),
+ StructField(name: "arr", dataType: .array(elementType: .string, containsNull: false)),
+ ])
+ #expect(structType.simpleString == "struct<id:bigint,arr:array<string>>")
+ #expect(StructType().simpleString == "struct<>")
+ }
+
+ @Test
+ func accessors() {
+ let structType = StructType(fields: [
+ StructField(name: "id", dataType: .long, nullable: false),
+ StructField(name: "name", dataType: .string),
+ StructField(name: "salary", dataType: .decimal(precision: 10, scale: 2)),
+ ])
+ #expect(structType.fieldNames == ["id", "name", "salary"])
+ #expect(structType.count == 3)
+ #expect(structType[0] == StructField(name: "id", dataType: .long, nullable: false))
+ #expect(structType["name"] == StructField(name: "name", dataType: .string))
+ #expect(structType["name"]?.nullable == true)
+ #expect(structType["nonexistent"] == nil)
+ #expect(structType.map { $0.dataType.simpleString } == ["bigint", "string", "decimal(10,2)"])
+
+ var count = 0
+ for field in structType {
+ #expect(structType[count] == field)
+ count += 1
+ }
+ #expect(count == 3)
+ }
+
+ @Test
+ func equatable() {
+ let a = StructType(fields: [StructField(name: "id", dataType: .long)])
+ let b = StructType(fields: [StructField(name: "id", dataType: .long)])
+ let c = StructType(fields: [StructField(name: "id", dataType: .integer)])
+ #expect(a == b)
+ #expect(a != c)
+ #expect(a != StructType())
+ #expect(
+ StructField(name: "id", dataType: .long)
+ != StructField(name: "id", dataType: .long, nullable: false))
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org