Skip to content

Commit c533d4c

Browse files
committed
[SPARK-52061] Add getOption/isModifiable and set variants to RuntimeConf
### What changes were proposed in this pull request? This PR aims to add `getOption`, `isModifiable` and more `set` API variants to `RuntimeConf`. ### Why are the changes needed? For feature parity. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Pass the CIs. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #119 from dongjoon-hyun/SPARK-52061. Authored-by: Dongjoon Hyun <[email protected]> Signed-off-by: Dongjoon Hyun <[email protected]>
1 parent dc2f57e commit c533d4c

File tree

3 files changed

+174
-1
lines changed

3 files changed

+174
-1
lines changed

Sources/SparkConnect/RuntimeConf.swift

Lines changed: 47 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,23 +35,69 @@ public actor RuntimeConf {
3535
try await client.setConf(map: [key: value])
3636
}
3737

38+
/// Set a new configuration.
39+
/// - Parameters:
40+
/// - key: A string for the configuration key.
41+
/// - value: A boolean value for the configuration.
42+
public func set(_ key: String, _ value: Bool) async throws {
43+
try await client.setConf(map: [key: String(value)])
44+
}
45+
46+
/// Set a new configuration.
47+
/// - Parameters:
48+
/// - key: A string for the configuration key.
49+
/// - value: A Int64 value for the configuration.
50+
public func set(_ key: String, _ value: Int64) async throws {
51+
try await client.setConf(map: [key: String(value)])
52+
}
53+
3854
/// Reset a configuration.
3955
/// - Parameters:
4056
/// - key: A string for the configuration key.
4157
public func unset(_ key: String) async throws {
4258
try await client.unsetConf(keys: [key])
4359
}
4460

45-
/// Get a configuration.
61+
/// Returns the value of Spark runtime configuration property for the given key. If the key is
62+
/// not set yet, return its default value if possible, otherwise `NoSuchElementException` will be
63+
/// thrown.
4664
/// - Parameter key: A string for the configuration look-up.
4765
/// - Returns: A string for the configuration.
4866
public func get(_ key: String) async throws -> String {
4967
return try await client.getConf(key)
5068
}
5169

70+
/// Returns the value of Spark runtime configuration property for the given key. If the key is
71+
/// not set yet, return the user given `value`. This is useful when its default value defined
72+
/// by Apache Spark is not the desired one.
73+
/// - Parameters:
74+
/// - key: A string for the configuration key.
75+
/// - value: A default string value for the configuration.
76+
public func get(_ key: String, _ value: String) async throws -> String {
77+
return try await client.getConfWithDefault(key, value)
78+
}
79+
5280
/// Get all configurations.
5381
/// - Returns: A map of configuration key-values.
5482
public func getAll() async throws -> [String: String] {
5583
return try await client.getConfAll()
5684
}
85+
86+
/// Returns the value of Spark runtime configuration property for the given key. If the key is
87+
/// not set yet, return its default value if possible, otherwise `nil` will be returned.
88+
/// - Parameter key: A string for the configuration look-up.
89+
/// - Returns: A string for the configuration or nil.
90+
public func getOption(_ key: String) async throws -> String? {
91+
return try await client.getConfOption(key)
92+
}
93+
94+
/// Indicates whether the configuration property with the given key is modifiable in the current
95+
/// session.
96+
/// - Parameter key: A string for the configuration look-up.
97+
/// - Returns: `true` if the configuration property is modifiable. For static SQL, Spark Core, invalid
98+
/// (not existing) and other non-modifiable configuration properties, the returned value is
99+
/// `false`.
100+
public func isModifiable(_ key: String) async throws -> Bool {
101+
return try await client.isModifiable(key)
102+
}
57103
}

Sources/SparkConnect/SparkConnectClient.swift

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -206,6 +206,73 @@ public actor SparkConnectClient {
206206
}
207207
}
208208

209+
/// Create a ``ConfigRequest`` instance for `GetWithDefault` operation.
210+
/// - Parameter pairs: A key-value dictionary.
211+
/// - Returns: A `ConfigRequest` instance.
212+
func getConfigRequestGetWithDefault(_ pairs: [String: String]) -> ConfigRequest {
213+
var request = ConfigRequest()
214+
request.operation = ConfigRequest.Operation()
215+
var getWithDefault = ConfigRequest.GetWithDefault()
216+
getWithDefault.pairs = pairs.toSparkConnectKeyValue
217+
request.operation.opType = .getWithDefault(getWithDefault)
218+
return request
219+
}
220+
221+
/// Returns the value of Spark runtime configuration property for the given key. If the key is
222+
/// not set yet, return the user given `value`. This is useful when its default value defined
223+
/// by Apache Spark is not the desired one.
224+
/// - Parameters:
225+
/// - key: A string for the configuration key.
226+
/// - value: A default value for the configuration.
227+
func getConfWithDefault(_ key: String, _ value: String) async throws -> String {
228+
try await withGPRC { client in
229+
let service = SparkConnectService.Client(wrapping: client)
230+
var request = getConfigRequestGetWithDefault([key: value])
231+
request.clientType = clientType
232+
request.userContext = userContext
233+
request.sessionID = self.sessionID!
234+
let response = try await service.config(request)
235+
let result = if response.pairs[0].hasValue {
236+
response.pairs[0].value
237+
} else {
238+
value
239+
}
240+
return result
241+
}
242+
}
243+
244+
/// Create a ``ConfigRequest`` instance for `GetOption` operation.
245+
/// - Parameter keys: An array of keys to get.
246+
/// - Returns: A `ConfigRequest` instance.
247+
func getConfigRequestGetOption(_ keys: [String]) -> ConfigRequest {
248+
var request = ConfigRequest()
249+
request.operation = ConfigRequest.Operation()
250+
var getOption = ConfigRequest.GetOption()
251+
getOption.keys = keys
252+
request.operation.opType = .getOption(getOption)
253+
return request
254+
}
255+
256+
/// Request the server to get a value of the given key.
257+
/// - Parameter key: A string for key to look up.
258+
/// - Returns: A string or nil for the value of the key.
259+
func getConfOption(_ key: String) async throws -> String? {
260+
try await withGPRC { client in
261+
let service = SparkConnectService.Client(wrapping: client)
262+
var request = getConfigRequestGetOption([key])
263+
request.clientType = clientType
264+
request.userContext = userContext
265+
request.sessionID = self.sessionID!
266+
let response = try await service.config(request)
267+
let result: String? = if response.pairs[0].hasValue {
268+
response.pairs[0].value
269+
} else {
270+
nil
271+
}
272+
return result
273+
}
274+
}
275+
209276
/// Create a ``ConfigRequest`` for `GetAll` operation.
210277
/// - Returns: A `ConfigRequest` instance.
211278
func getConfigRequestGetAll() -> ConfigRequest {
@@ -234,6 +301,35 @@ public actor SparkConnectClient {
234301
}
235302
}
236303

304+
/// Create a ``ConfigRequest`` for `IsModifiable` operation.
305+
/// - Returns: A `ConfigRequest` instance.
306+
func getConfigRequestIsModifiable(_ keys: [String]) -> ConfigRequest {
307+
var request = ConfigRequest()
308+
request.operation = ConfigRequest.Operation()
309+
var isModifiable = ConfigRequest.IsModifiable()
310+
isModifiable.keys = keys
311+
request.operation.opType = .isModifiable(isModifiable)
312+
return request
313+
}
314+
315+
/// Indicates whether the configuration property with the given key is modifiable in the current
316+
/// session.
317+
/// - Parameter key: A string for the configuration look-up.
318+
/// - Returns: `true` if the configuration property is modifiable. For static SQL, Spark Core, invalid
319+
/// (not existing) and other non-modifiable configuration properties, the returned value is
320+
/// `false`.
321+
func isModifiable(_ key: String) async throws -> Bool {
322+
try await withGPRC { client in
323+
let service = SparkConnectService.Client(wrapping: client)
324+
var request = getConfigRequestIsModifiable([key])
325+
request.clientType = clientType
326+
request.userContext = userContext
327+
request.sessionID = self.sessionID!
328+
let response = try await service.config(request)
329+
return response.pairs[0].value == "true"
330+
}
331+
}
332+
237333
func getLocalRelation() -> Plan {
238334
var localRelation = Spark_Connect_LocalRelation()
239335
localRelation.schema = ""

Tests/SparkConnectTests/RuntimeConfTests.swift

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,26 @@ struct RuntimeConfTests {
4242
await client.stop()
4343
}
4444

45+
@Test
46+
func getWithDefault() async throws {
47+
let client = SparkConnectClient(remote: TEST_REMOTE)
48+
try await client.connect(UUID().uuidString)
49+
let conf = RuntimeConf(client)
50+
#expect(try await conf.get("spark.sql.adaptive.customCostEvaluatorClass", "XYZ") == "XYZ")
51+
#expect(try await conf.get("spark.test.non-exist", "my_default") == "my_default")
52+
await client.stop()
53+
}
54+
55+
@Test
56+
func getOption() async throws {
57+
let client = SparkConnectClient(remote: TEST_REMOTE)
58+
try await client.connect(UUID().uuidString)
59+
let conf = RuntimeConf(client)
60+
#expect(try await conf.getOption("spark.app.name") != nil)
61+
#expect(try await conf.getOption("spark.test.non-exist") == nil)
62+
await client.stop()
63+
}
64+
4565
@Test
4666
func set() async throws {
4767
let client = SparkConnectClient(remote: TEST_REMOTE)
@@ -86,4 +106,15 @@ struct RuntimeConfTests {
86106
#expect(map["spark.master"] != nil)
87107
await client.stop()
88108
}
109+
110+
@Test
111+
func isModifiable() async throws {
112+
let client = SparkConnectClient(remote: TEST_REMOTE)
113+
try await client.connect(UUID().uuidString)
114+
let conf = RuntimeConf(client)
115+
#expect(try await conf.isModifiable("spark.sql.adaptive.customCostEvaluatorClass"))
116+
#expect(try await conf.isModifiable("spark.sql.warehouse.dir") == false)
117+
#expect(try await conf.isModifiable("spark.test.non-exist") == false)
118+
await client.stop()
119+
}
89120
}

0 commit comments

Comments
 (0)