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/10/20 21:50:46 UTC
(spark-connect-swift) branch main updated: [SPARK-60407] Fix command APIs to not return empty results when ExecutionProgress arrives first
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 585f7a3 [SPARK-60407] Fix command APIs to not return empty results when ExecutionProgress arrives first
585f7a3 is described below
commit 585f7a37b3b4ab9a90ebeaa188c81f14667a65b7
Author: Dongjoon Hyun <do...@apache.org>
AuthorDate: Tue Oct 20 14:50:40 2026 -0700
[SPARK-60407] Fix command APIs to not return empty results when ExecutionProgress arrives first
### What changes were proposed in this pull request?
This PR fixes the command APIs to read the result from the first response of the expected `responseType` instead of the first response.
- Add `[ExecutePlanResponse].firstOrThrow(_:)`, which takes the key path of a `responseType` case like `\.streamingQueryCommandResult` and throws `SparkConnectError.invalidState` if no response has it.
- Use it at all 14 call sites that read a command result: `StreamingQuery` (7), `StreamingQueryManager` (3), `SparkConnectClient` (3, for `SparkSession.executeCommand`, `DataFrame.checkpoint`/`localCheckpoint`, and `createDataflowGraph`), and `DataStreamWriter.start` (1).
The other 17 `firstOrThrow()` call sites in `Catalog` and `DataFrameStatFunctions` are unchanged because they work on the `[Row]` of `collect()`, which only reads the schema and Arrow batches.
### Why are the changes needed?
Since Apache Spark 4.0.0 (SPARK-47081), the Spark Connect server sends an `ExecutionProgress` response every `spark.connect.progress.reportInterval` (default: `2s`) while an execution is waiting for its next response. So, a command taking longer than 2 seconds receives `ExecutionProgress` responses before its result. For example, `apache/spark:4.2.0` returns the following.
```
checkpoint : executionProgress (x3), checkpointCommandResult, executionProgress (x2), resultComplete
awaitTermination : executionProgress, streamingQueryCommandResult, executionProgress, resultComplete
```
Since the client used the first response, it silently returned the empty default value of the expected result.
- `StreamingQuery.awaitTermination(timeoutMs)` and `StreamingQueryManager.awaitAnyTermination(timeoutMs)` return `false` even when the query terminated.
- Eager `DataFrame.checkpoint()` and `localCheckpoint()` return a `DataFrame` with an empty cached relation ID, which fails at the next action with `[CONNECT_INVALID_PLAN.DATAFRAME_NOT_FOUND]`.
The reference clients handle this already. The Scala client ignores `ExecutionProgress` in `SparkSession.execute(command)` and uses `find(_.hasSqlCommandResult)`, and the Python client handles each response by `HasField`.
### Does this PR introduce _any_ user-facing change?
Yes. Commands taking longer than `spark.connect.progress.reportInterval` return the actual results instead of empty default values. If the server returns no expected result, `SparkConnectError.invalidState` is thrown instead.
### How was this patch tested?
Pass the CIs with the newly added `ExtensionTests` cases, which need no Spark Connect server.
- `firstOrThrowCommandResultAfterExecutionProgress`
- `firstOrThrowEmptyCommandResult`
- `firstOrThrowWithoutCommandResult`
Also, manually verified against an `apache/spark:4.2.0` Connect server that eager `localCheckpoint()` of `SELECT reflect('java.lang.Thread', 'sleep', 3000L)` works and `awaitTermination(10000)` returns `true` for a query stopped after 3 seconds.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 5.5
Closes #609 from dongjoon-hyun/SPARK-60407.
Authored-by: Dongjoon Hyun <do...@apache.org>
Signed-off-by: Dongjoon Hyun <do...@apache.org>
---
Sources/SparkConnect/DataStreamWriter.swift | 2 +-
Sources/SparkConnect/Extension.swift | 27 +++++++++++++++
Sources/SparkConnect/SparkConnectClient.swift | 6 ++--
Sources/SparkConnect/StreamingQuery.swift | 14 ++++----
Sources/SparkConnect/StreamingQueryManager.swift | 6 ++--
Tests/SparkConnectTests/ExtensionTests.swift | 42 ++++++++++++++++++++++++
6 files changed, 83 insertions(+), 14 deletions(-)
diff --git a/Sources/SparkConnect/DataStreamWriter.swift b/Sources/SparkConnect/DataStreamWriter.swift
index f33f6c4..36f08bf 100644
--- a/Sources/SparkConnect/DataStreamWriter.swift
+++ b/Sources/SparkConnect/DataStreamWriter.swift
@@ -185,7 +185,7 @@ public actor DataStreamWriter: Sendable {
command.writeStreamOperationStart = writeStreamOperationStart
let response = try await df.spark.client.execute(df.spark.sessionID, command)
- let result = try response.firstOrThrow().writeStreamOperationStartResult
+ let result = try response.firstOrThrow(\.writeStreamOperationStartResult)
// `result.queryStartedEventJson` is intentionally ignored because this client does not
// support `StreamingQueryListener` yet. Reference clients (PySpark/Scala) parse it and
// post it to their listener bus, which is a no-op when no listener is registered.
diff --git a/Sources/SparkConnect/Extension.swift b/Sources/SparkConnect/Extension.swift
index dc0575f..b369ac3 100644
--- a/Sources/SparkConnect/Extension.swift
+++ b/Sources/SparkConnect/Extension.swift
@@ -273,6 +273,33 @@ extension Array {
}
}
+extension [ExecutePlanResponse] {
+ /// Get the `result` of the first response whose `responseType` is the case of `result`, or throw
+ /// ``SparkConnectError/invalidState(_:)`` if there is none.
+ ///
+ /// The Spark Connect server can send other responses like `ExecutionProgress` before the result
+ /// of a long-running command, so the result is not always the first response.
+ /// - Parameters:
+ /// - result: The key path of a `responseType` case like `\.streamingQueryCommandResult`.
+ /// - function: The caller's name which is used in the error message.
+ /// - Returns: The result of the first response of the `result` case.
+ func firstOrThrow<Value>(
+ _ result: WritableKeyPath<ExecutePlanResponse, Value>, _ function: String = #function
+ ) throws -> Value {
+ for response in self {
+ // Setting `result` of an empty response selects the `responseType` case of `result`.
+ var expected = ExecutePlanResponse()
+ expected[keyPath: result] = response[keyPath: result]
+ if expected.responseType == response.responseType {
+ return response[keyPath: result]
+ }
+ }
+ throw SparkConnectError.invalidState(
+ SparkConnectError.Details(
+ message: "The server returned no \(Value.self) for \(function)."))
+ }
+}
+
extension Data {
/// Get an `Int32` value from unsafe 4 bytes.
var int32: Int32 { withUnsafeBytes({ $0.load(as: Int32.self) }) }
diff --git a/Sources/SparkConnect/SparkConnectClient.swift b/Sources/SparkConnect/SparkConnectClient.swift
index eb4e8ef..050f449 100644
--- a/Sources/SparkConnect/SparkConnectClient.swift
+++ b/Sources/SparkConnect/SparkConnectClient.swift
@@ -972,7 +972,7 @@ public actor SparkConnectClient {
var command = Command()
command.commandType = .executeExternalCommand(executeExternalCommand)
let response = try await execute(self.sessionID!, command)
- let relation = try response.firstOrThrow().sqlCommandResult.relation
+ let relation = try response.firstOrThrow(\.sqlCommandResult).relation
var plan = Plan()
plan.opType = .root(relation)
return plan
@@ -1397,7 +1397,7 @@ public actor SparkConnectClient {
var command = Spark_Connect_Command()
command.checkpointCommand = checkpointCommand
let response = try await execute(self.sessionID!, command)
- let cachedRemoteRelation = try response.firstOrThrow().checkpointCommandResult.relation
+ let cachedRemoteRelation = try response.firstOrThrow(\.checkpointCommandResult).relation
return Self.createPlan { $0.cachedRemoteRelation = cachedRemoteRelation }
}
@@ -1488,7 +1488,7 @@ public actor SparkConnectClient {
command.commandType = .pipelineCommand(pipelineCommand)
let response = try await execute(self.sessionID!, command)
- let result = try response.firstOrThrow().pipelineCommandResult.createDataflowGraphResult
+ let result = try response.firstOrThrow(\.pipelineCommandResult).createDataflowGraphResult
return result.dataflowGraphID
}
diff --git a/Sources/SparkConnect/StreamingQuery.swift b/Sources/SparkConnect/StreamingQuery.swift
index c3bb36b..bff28fe 100644
--- a/Sources/SparkConnect/StreamingQuery.swift
+++ b/Sources/SparkConnect/StreamingQuery.swift
@@ -98,7 +98,7 @@ public actor StreamingQuery: Sendable {
public var isActive: Bool {
get async throws {
let response = try await executeCommand(StreamingQueryCommand.OneOf_Command.status(true))
- return try response.firstOrThrow().streamingQueryCommandResult.status.isActive
+ return try response.firstOrThrow(\.streamingQueryCommandResult).status.isActive
}
}
@@ -106,7 +106,7 @@ public actor StreamingQuery: Sendable {
/// - Returns: A ``StreamingQueryException`` if the query was terminated by an exception, `nil` otherwise.
public func exception() async throws -> StreamingQueryException? {
let response = try await executeCommand(StreamingQueryCommand.OneOf_Command.exception(true))
- let result = try response.firstOrThrow().streamingQueryCommandResult.exception
+ let result = try response.firstOrThrow(\.streamingQueryCommandResult).exception
guard result.hasExceptionMessage else {
return nil
}
@@ -121,7 +121,7 @@ public actor StreamingQuery: Sendable {
/// - Returns:
public func status() async throws -> StreamingQueryStatus {
let response = try await executeCommand(StreamingQueryCommand.OneOf_Command.status(true))
- let result = try response.firstOrThrow().streamingQueryCommandResult.status
+ let result = try response.firstOrThrow(\.streamingQueryCommandResult).status
return StreamingQueryStatus(
statusMessage: result.statusMessage,
isDataAvailable: result.isDataAvailable,
@@ -137,7 +137,7 @@ public actor StreamingQuery: Sendable {
get async throws {
let response = try await executeCommand(
StreamingQueryCommand.OneOf_Command.recentProgress(true))
- let result = try response.firstOrThrow().streamingQueryCommandResult.recentProgress
+ let result = try response.firstOrThrow(\.streamingQueryCommandResult).recentProgress
return try result.recentProgressJson.map { try StreamingQueryProgress.fromJson($0) }
}
}
@@ -147,7 +147,7 @@ public actor StreamingQuery: Sendable {
get async throws {
let response = try await executeCommand(
StreamingQueryCommand.OneOf_Command.lastProgress(true))
- let result = try response.firstOrThrow().streamingQueryCommandResult.recentProgress
+ let result = try response.firstOrThrow(\.streamingQueryCommandResult).recentProgress
guard let json = result.recentProgressJson.first else {
return nil
}
@@ -170,7 +170,7 @@ public actor StreamingQuery: Sendable {
}
let response = try await executeCommand(
StreamingQueryCommand.OneOf_Command.awaitTermination(command))
- return try response.firstOrThrow().streamingQueryCommandResult.awaitTermination.terminated
+ return try response.firstOrThrow(\.streamingQueryCommandResult).awaitTermination.terminated
}
/// Blocks until all available data in the source has been processed and committed to the sink.
@@ -201,6 +201,6 @@ public actor StreamingQuery: Sendable {
var command = Spark_Connect_StreamingQueryCommand.ExplainCommand()
command.extended = extended
let response = try await executeCommand(StreamingQueryCommand.OneOf_Command.explain(command))
- print(try response.firstOrThrow().streamingQueryCommandResult.explain.result)
+ print(try response.firstOrThrow(\.streamingQueryCommandResult).explain.result)
}
}
diff --git a/Sources/SparkConnect/StreamingQueryManager.swift b/Sources/SparkConnect/StreamingQueryManager.swift
index 44923f1..c6dafc7 100644
--- a/Sources/SparkConnect/StreamingQueryManager.swift
+++ b/Sources/SparkConnect/StreamingQueryManager.swift
@@ -35,7 +35,7 @@ public actor StreamingQueryManager {
get async throws {
let command = StreamingQueryManagerCommand.OneOf_Command.active(true)
let response = try await self.sparkSession.client.executeStreamingQueryManagerCommand(command)
- let result = try response.firstOrThrow().streamingQueryManagerCommandResult
+ let result = try response.firstOrThrow(\.streamingQueryManagerCommandResult)
return try result.active.activeQueries.map {
try StreamingQuery($0.id, $0.name, self.sparkSession)
}
@@ -55,7 +55,7 @@ public actor StreamingQueryManager {
public func get(_ id: String) async throws -> StreamingQuery {
let command = StreamingQueryManagerCommand.OneOf_Command.getQuery(id)
let response = try await self.sparkSession.client.executeStreamingQueryManagerCommand(command)
- let query = try response.firstOrThrow().streamingQueryManagerCommandResult.query
+ let query = try response.firstOrThrow(\.streamingQueryManagerCommandResult).query
guard query.hasID else {
throw SparkConnectError.InvalidArgument
}
@@ -78,7 +78,7 @@ public actor StreamingQueryManager {
let command = StreamingQueryManagerCommand.OneOf_Command.awaitAnyTermination(
awaitAnyTerminationCommand)
let response = try await self.sparkSession.client.executeStreamingQueryManagerCommand(command)
- let result = try response.firstOrThrow().streamingQueryManagerCommandResult
+ let result = try response.firstOrThrow(\.streamingQueryManagerCommandResult)
return result.awaitAnyTermination.terminated
}
diff --git a/Tests/SparkConnectTests/ExtensionTests.swift b/Tests/SparkConnectTests/ExtensionTests.swift
index 098c66f..38b6339 100644
--- a/Tests/SparkConnectTests/ExtensionTests.swift
+++ b/Tests/SparkConnectTests/ExtensionTests.swift
@@ -38,4 +38,46 @@ struct ExtensionTests {
#expect(
error.message == "The server returned an empty response for firstOrThrowOnEmptyArray().")
}
+
+ @Test
+ func firstOrThrowCommandResultAfterExecutionProgress() throws {
+ // The server sends an empty `ExecutionProgress` periodically while a command is running.
+ var progress = ExecutePlanResponse()
+ progress.executionProgress = ExecutePlanResponse.ExecutionProgress()
+ var result = ExecutePlanResponse()
+ result.streamingQueryCommandResult.awaitTermination.terminated = true
+ var complete = ExecutePlanResponse()
+ complete.resultComplete = ExecutePlanResponse.ResultComplete()
+ let responses = [progress, progress, result, complete]
+ #expect(try responses.firstOrThrow(\.streamingQueryCommandResult).awaitTermination.terminated)
+ }
+
+ @Test
+ func firstOrThrowEmptyCommandResult() throws {
+ // An empty result is found by its case while the responses of no or other cases are skipped.
+ var metrics = ExecutePlanResponse()
+ metrics.metrics = ExecutePlanResponse.Metrics()
+ var other = ExecutePlanResponse()
+ other.sqlCommandResult = ExecutePlanResponse.SqlCommandResult()
+ var result = ExecutePlanResponse()
+ result.checkpointCommandResult = Spark_Connect_CheckpointCommandResult()
+ let responses = [metrics, other, result]
+ #expect(try responses.firstOrThrow(\.checkpointCommandResult) == result.checkpointCommandResult)
+ }
+
+ @Test
+ func firstOrThrowWithoutCommandResult() throws {
+ var progress = ExecutePlanResponse()
+ progress.executionProgress = ExecutePlanResponse.ExecutionProgress()
+ var complete = ExecutePlanResponse()
+ complete.resultComplete = ExecutePlanResponse.ResultComplete()
+ let error = try #require(throws: SparkConnectError.self) {
+ _ = try [progress, complete].firstOrThrow(\.checkpointCommandResult)
+ }
+ #expect(error == .InvalidState)
+ #expect(
+ error.message
+ == "The server returned no Spark_Connect_CheckpointCommandResult for "
+ + "firstOrThrowWithoutCommandResult().")
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org