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 2025/04/20 21:56:34 UTC
(spark-connect-swift) branch main updated: [SPARK-51851] Refactor to use `withGPRC` wrappers
This is an automated email from the ASF dual-hosted git repository.
dongjoon 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 1169bea [SPARK-51851] Refactor to use `withGPRC` wrappers
1169bea is described below
commit 1169beab767b7a41801349c9b81647eadea2e1f0
Author: Dongjoon Hyun <do...@apache.org>
AuthorDate: Mon Apr 21 06:55:07 2025 +0900
[SPARK-51851] Refactor to use `withGPRC` wrappers
### What changes were proposed in this pull request?
This PR aims to refactor to use `withGRPC` wrappers in `SparkConnectClient` and `DataFrame`.
### Why are the changes needed?
This is helpful not only for reducing the duplicated code but also introducing a security feature.
### Does this PR introduce _any_ user-facing change?
No, this is an internal code refactoring.
### How was this patch tested?
Pass the CIs.
### Was this patch authored or co-authored using generative AI tooling?
No.
This patch had conflicts when merged, resolved by
Committer: Dongjoon Hyun <do...@apache.org>
Closes #73 from dongjoon-hyun/SPARK-51851.
Authored-by: Dongjoon Hyun <do...@apache.org>
Signed-off-by: Dongjoon Hyun <do...@apache.org>
---
Sources/SparkConnect/DataFrame.swift | 86 +++++++--------------------
Sources/SparkConnect/SparkConnectClient.swift | 69 +++++++--------------
2 files changed, 43 insertions(+), 112 deletions(-)
diff --git a/Sources/SparkConnect/DataFrame.swift b/Sources/SparkConnect/DataFrame.swift
index d3b9abd..42f4be9 100644
--- a/Sources/SparkConnect/DataFrame.swift
+++ b/Sources/SparkConnect/DataFrame.swift
@@ -110,16 +110,24 @@ public actor DataFrame: Sendable {
}
}
- private func analyzePlanIfNeeded() async throws {
- if self._schema != nil {
- return
- }
+ private func withGPRC<Result: Sendable>(
+ _ f: (GRPCClient<GRPCNIOTransportHTTP2.HTTP2ClientTransport.Posix>) async throws -> Result
+ ) async throws -> Result {
try await withGRPCClient(
transport: .http2NIOPosix(
target: .dns(host: spark.client.host, port: spark.client.port),
transportSecurity: .plaintext
)
) { client in
+ return try await f(client)
+ }
+ }
+
+ private func analyzePlanIfNeeded() async throws {
+ if self._schema != nil {
+ return
+ }
+ try await withGPRC { client in
let service = Spark_Connect_SparkConnectService.Client(wrapping: client)
let response = try await service.analyzePlan(
spark.client.getAnalyzePlanRequest(spark.sessionID, plan))
@@ -132,12 +140,7 @@ public actor DataFrame: Sendable {
public func count() async throws -> Int64 {
let counter = Atomic(Int64(0))
- try await withGRPCClient(
- transport: .http2NIOPosix(
- target: .dns(host: spark.client.host, port: spark.client.port),
- transportSecurity: .plaintext
- )
- ) { client in
+ try await withGPRC { client in
let service = Spark_Connect_SparkConnectService.Client(wrapping: client)
try await service.executePlan(spark.client.getExecutePlanRequest(plan)) {
response in
@@ -154,12 +157,7 @@ public actor DataFrame: Sendable {
// Clear all existing batches.
self.batches.removeAll()
- try await withGRPCClient(
- transport: .http2NIOPosix(
- target: .dns(host: spark.client.host, port: spark.client.port),
- transportSecurity: .plaintext
- )
- ) { client in
+ try await withGPRC { client in
let service = Spark_Connect_SparkConnectService.Client(wrapping: client)
try await service.executePlan(spark.client.getExecutePlanRequest(plan)) {
response in
@@ -397,12 +395,7 @@ public actor DataFrame: Sendable {
/// (without any Spark executors).
/// - Returns: True if the plan is local.
public func isLocal() async throws -> Bool {
- try await withGRPCClient(
- transport: .http2NIOPosix(
- target: .dns(host: spark.client.host, port: spark.client.port),
- transportSecurity: .plaintext
- )
- ) { client in
+ try await withGPRC { client in
let service = Spark_Connect_SparkConnectService.Client(wrapping: client)
let response = try await service.analyzePlan(spark.client.getIsLocal(spark.sessionID, plan))
return response.isLocal.isLocal
@@ -413,12 +406,7 @@ public actor DataFrame: Sendable {
/// arrives.
/// - Returns: True if a plan is streaming.
public func isStreaming() async throws -> Bool {
- try await withGRPCClient(
- transport: .http2NIOPosix(
- target: .dns(host: spark.client.host, port: spark.client.port),
- transportSecurity: .plaintext
- )
- ) { client in
+ try await withGPRC { client in
let service = Spark_Connect_SparkConnectService.Client(wrapping: client)
let response = try await service.analyzePlan(spark.client.getIsStreaming(spark.sessionID, plan))
return response.isStreaming.isStreaming
@@ -442,12 +430,7 @@ public actor DataFrame: Sendable {
public func persist(storageLevel: StorageLevel = StorageLevel.MEMORY_AND_DISK) async throws
-> DataFrame
{
- try await withGRPCClient(
- transport: .http2NIOPosix(
- target: .dns(host: spark.client.host, port: spark.client.port),
- transportSecurity: .plaintext
- )
- ) { client in
+ try await withGPRC { client in
let service = Spark_Connect_SparkConnectService.Client(wrapping: client)
_ = try await service.analyzePlan(
spark.client.getPersist(spark.sessionID, plan, storageLevel))
@@ -461,12 +444,7 @@ public actor DataFrame: Sendable {
/// - Parameter blocking: Whether to block until all blocks are deleted.
/// - Returns: A `DataFrame`
public func unpersist(blocking: Bool = false) async throws -> DataFrame {
- try await withGRPCClient(
- transport: .http2NIOPosix(
- target: .dns(host: spark.client.host, port: spark.client.port),
- transportSecurity: .plaintext
- )
- ) { client in
+ try await withGPRC { client in
let service = Spark_Connect_SparkConnectService.Client(wrapping: client)
_ = try await service.analyzePlan(spark.client.getUnpersist(spark.sessionID, plan, blocking))
}
@@ -476,12 +454,7 @@ public actor DataFrame: Sendable {
public var storageLevel: StorageLevel {
get async throws {
- try await withGRPCClient(
- transport: .http2NIOPosix(
- target: .dns(host: spark.client.host, port: spark.client.port),
- transportSecurity: .plaintext
- )
- ) { client in
+ try await withGPRC { client in
let service = Spark_Connect_SparkConnectService.Client(wrapping: client)
return try await service
.analyzePlan(spark.client.getStorageLevel(spark.sessionID, plan)).getStorageLevel.storageLevel.toStorageLevel
@@ -508,12 +481,7 @@ public actor DataFrame: Sendable {
/// - Parameter mode: the expected output format of plans;
/// `simple`, `extended`, `codegen`, `cost`, `formatted`.
public func explain(_ mode: String) async throws {
- try await withGRPCClient(
- transport: .http2NIOPosix(
- target: .dns(host: spark.client.host, port: spark.client.port),
- transportSecurity: .plaintext
- )
- ) { client in
+ try await withGPRC { client in
let service = Spark_Connect_SparkConnectService.Client(wrapping: client)
let response = try await service.analyzePlan(spark.client.getExplain(spark.sessionID, plan, mode))
print(response.explain.explainString)
@@ -525,12 +493,7 @@ public actor DataFrame: Sendable {
/// results. Depending on the source relations, this may not find all input files. Duplicates are removed.
/// - Returns: An array of file path strings.
public func inputFiles() async throws -> [String] {
- try await withGRPCClient(
- transport: .http2NIOPosix(
- target: .dns(host: spark.client.host, port: spark.client.port),
- transportSecurity: .plaintext
- )
- ) { client in
+ try await withGPRC { client in
let service = Spark_Connect_SparkConnectService.Client(wrapping: client)
let response = try await service.analyzePlan(spark.client.getInputFiles(spark.sessionID, plan))
return response.inputFiles.files
@@ -545,12 +508,7 @@ public actor DataFrame: Sendable {
/// Prints the schema up to the given level to the console in a nice tree format.
/// - Parameter level: A level to be printed.
public func printSchema(_ level: Int32) async throws {
- try await withGRPCClient(
- transport: .http2NIOPosix(
- target: .dns(host: spark.client.host, port: spark.client.port),
- transportSecurity: .plaintext
- )
- ) { client in
+ try await withGPRC { client in
let service = Spark_Connect_SparkConnectService.Client(wrapping: client)
let response = try await service.analyzePlan(spark.client.getTreeString(spark.sessionID, plan, level))
print(response.treeString.treeString)
diff --git a/Sources/SparkConnect/SparkConnectClient.swift b/Sources/SparkConnect/SparkConnectClient.swift
index e058177..4854159 100644
--- a/Sources/SparkConnect/SparkConnectClient.swift
+++ b/Sources/SparkConnect/SparkConnectClient.swift
@@ -52,12 +52,7 @@ public actor SparkConnectClient {
/// - Parameter sessionID: A string for the session ID.
/// - Returns: An `AnalyzePlanResponse` instance for `SparkVersion`
func connect(_ sessionID: String) async throws -> AnalyzePlanResponse {
- try await withGRPCClient(
- transport: .http2NIOPosix(
- target: .dns(host: self.host, port: self.port),
- transportSecurity: .plaintext
- )
- ) { client in
+ try await withGPRC { client in
// To prevent server-side `INVALID_HANDLE.FORMAT (SQLSTATE: HY000)` exception.
if UUID(uuidString: sessionID) == nil {
throw SparkConnectError.InvalidSessionIDException
@@ -73,6 +68,19 @@ public actor SparkConnectClient {
}
}
+ private func withGPRC<Result: Sendable>(
+ _ f: (GRPCClient<GRPCNIOTransportHTTP2.HTTP2ClientTransport.Posix>) async throws -> Result
+ ) async throws -> Result {
+ try await withGRPCClient(
+ transport: .http2NIOPosix(
+ target: .dns(host: self.host, port: self.port),
+ transportSecurity: .plaintext
+ )
+ ) { client in
+ return try await f(client)
+ }
+ }
+
/// Create a ``ConfigRequest`` instance for `Set` operation.
/// - Parameter map: A map of key-value string pairs.
/// - Returns: A ``ConfigRequest`` instance.
@@ -89,12 +97,7 @@ public actor SparkConnectClient {
/// - Parameter map: A map of key-value pairs to set.
/// - Returns: Always return true.
func setConf(map: [String: String]) async throws -> Bool {
- try await withGRPCClient(
- transport: .http2NIOPosix(
- target: .dns(host: self.host, port: self.port),
- transportSecurity: .plaintext
- )
- ) { client in
+ try await withGPRC { client in
let service = SparkConnectService.Client(wrapping: client)
var request = getConfigRequestSet(map: map)
request.clientType = clientType
@@ -118,12 +121,7 @@ public actor SparkConnectClient {
}
func unsetConf(keys: [String]) async throws -> Bool {
- try await withGRPCClient(
- transport: .http2NIOPosix(
- target: .dns(host: self.host, port: self.port),
- transportSecurity: .plaintext
- )
- ) { client in
+ try await withGPRC { client in
let service = SparkConnectService.Client(wrapping: client)
var request = getConfigRequestUnset(keys: keys)
request.clientType = clientType
@@ -150,12 +148,7 @@ public actor SparkConnectClient {
/// - Parameter key: A string for key to look up.
/// - Returns: A string for the value of the key.
func getConf(_ key: String) async throws -> String {
- try await withGRPCClient(
- transport: .http2NIOPosix(
- target: .dns(host: self.host, port: self.port),
- transportSecurity: .plaintext
- )
- ) { client in
+ try await withGPRC { client in
let service = SparkConnectService.Client(wrapping: client)
var request = getConfigRequestGet(keys: [key])
request.clientType = clientType
@@ -179,12 +172,7 @@ public actor SparkConnectClient {
/// Request the server to get all configurations.
/// - Returns: A map of key-value pairs.
func getConfAll() async throws -> [String: String] {
- try await withGRPCClient(
- transport: .http2NIOPosix(
- target: .dns(host: self.host, port: self.port),
- transportSecurity: .plaintext
- )
- ) { client in
+ try await withGPRC { client in
let service = SparkConnectService.Client(wrapping: client)
var request = getConfigRequestGetAll()
request.clientType = clientType
@@ -451,12 +439,7 @@ public actor SparkConnectClient {
func execute(_ sessionID: String, _ command: Command) async throws -> [ExecutePlanResponse] {
self.result.removeAll()
- try await withGRPCClient(
- transport: .http2NIOPosix(
- target: .dns(host: self.host, port: self.port),
- transportSecurity: .plaintext
- )
- ) { client in
+ try await withGPRC { client in
let service = SparkConnectService.Client(wrapping: client)
var plan = Plan()
plan.opType = .command(command)
@@ -501,12 +484,7 @@ public actor SparkConnectClient {
/// - Parameter ddlString: A string to parse.
/// - Returns: A ``Spark_Connect_DataType`` instance.
func ddlParse(_ ddlString: String) async throws -> Spark_Connect_DataType {
- try await withGRPCClient(
- transport: .http2NIOPosix(
- target: .dns(host: self.host, port: self.port),
- transportSecurity: .plaintext
- )
- ) { client in
+ try await withGPRC { client in
let service = SparkConnectService.Client(wrapping: client)
let request = analyze(self.sessionID!, {
var ddlParse = AnalyzePlanRequest.DDLParse()
@@ -522,12 +500,7 @@ public actor SparkConnectClient {
/// - Parameter jsonString: A JSON string.
/// - Returns: A DDL string.
func jsonToDdl(_ jsonString: String) async throws -> String {
- try await withGRPCClient(
- transport: .http2NIOPosix(
- target: .dns(host: self.host, port: self.port),
- transportSecurity: .plaintext
- )
- ) { client in
+ try await withGPRC { client in
let service = SparkConnectService.Client(wrapping: client)
let request = analyze(self.sessionID!, {
var jsonToDDL = AnalyzePlanRequest.JsonToDDL()
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org