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/03/07 20:30:26 UTC
(spark) branch master updated: [SPARK-51425][CONNECT] Add client API to set custom `operation_id`
This is an automated email from the ASF dual-hosted git repository.
dongjoon pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 50e00b7b72a8 [SPARK-51425][CONNECT] Add client API to set custom `operation_id`
50e00b7b72a8 is described below
commit 50e00b7b72a8d3d9b103c5ea6f204a3e1e365eaf
Author: vicennial <ve...@databricks.com>
AuthorDate: Fri Mar 7 12:30:12 2025 -0800
[SPARK-51425][CONNECT] Add client API to set custom `operation_id`
### What changes were proposed in this pull request?
Adds an additional optional parameter to the Scala/Python APIs to allow a user to explicitly set an `operation_id`.
### Why are the changes needed?
The Spark Connect [protocol](https://github.com/apache/spark/blob/44e751f243a5a7be8e64c306e40ddc1502f72710/sql/connect/common/src/main/protobuf/spark/connect/base.proto#L318) allows the client to set an optional operation ID. However, there is no API that lets a user set this explicitly (although the client does set it in the case of Reaatchable Execution).
### Does this PR introduce _any_ user-facing change?
Yes.
Scala usage:
```scala
client.execute(plan, operationId = Some("10a4c38e-7e87-40ee-9d6f-60ff0751e63b))
```
Python usage:
```python
req = client._execute_plan_request_with_metadata(operation_id="10a4c38e-7e87-40ee-9d6f-60ff0751e63b")
# continue using the req as usual
```
### How was this patch tested?
New unit tests
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #50191 from vicennial/customOpId.
Authored-by: vicennial <ve...@databricks.com>
Signed-off-by: Dongjoon Hyun <do...@apache.org>
---
python/pyspark/sql/connect/client/core.py | 13 ++++++++++++-
.../pyspark/sql/tests/connect/client/test_client.py | 11 +++++++++++
.../sql/connect/client/SparkConnectClientSuite.scala | 20 ++++++++++++++++++++
.../sql/connect/client/SparkConnectClient.scala | 20 +++++++++++++++++++-
4 files changed, 62 insertions(+), 2 deletions(-)
diff --git a/python/pyspark/sql/connect/client/core.py b/python/pyspark/sql/connect/client/core.py
index 5327b97601d7..a672968af393 100644
--- a/python/pyspark/sql/connect/client/core.py
+++ b/python/pyspark/sql/connect/client/core.py
@@ -1228,7 +1228,9 @@ class SparkConnectClient(object):
"""
return self._builder.token
- def _execute_plan_request_with_metadata(self) -> pb2.ExecutePlanRequest:
+ def _execute_plan_request_with_metadata(
+ self, operation_id: Optional[str] = None
+ ) -> pb2.ExecutePlanRequest:
req = pb2.ExecutePlanRequest(
session_id=self._session_id,
client_type=self._builder.userAgent,
@@ -1238,6 +1240,15 @@ class SparkConnectClient(object):
req.client_observed_server_side_session_id = self._server_session_id
if self._user_id:
req.user_context.user_id = self._user_id
+ if operation_id is not None:
+ try:
+ uuid.UUID(operation_id, version=4)
+ except ValueError as ve:
+ raise PySparkValueError(
+ errorClass="INVALID_OPERATION_UUID_ID",
+ messageParameters={"arg_name": "operation_id", "origin": str(ve)},
+ )
+ req.operation_id = operation_id
return req
def _analyze_plan_request_with_metadata(self) -> pb2.AnalyzePlanRequest:
diff --git a/python/pyspark/sql/tests/connect/client/test_client.py b/python/pyspark/sql/tests/connect/client/test_client.py
index 741d6b9c1104..cba8bf3e1477 100644
--- a/python/pyspark/sql/tests/connect/client/test_client.py
+++ b/python/pyspark/sql/tests/connect/client/test_client.py
@@ -137,6 +137,7 @@ if should_test_connect:
self.req = req
resp = proto.ExecutePlanResponse()
resp.session_id = self._session_id
+ resp.operation_id = req.operation_id
pdf = pd.DataFrame(data={"col1": [1, 2]})
schema = pa.Schema.from_pandas(pdf)
@@ -255,6 +256,16 @@ class SparkConnectClientTestCase(unittest.TestCase):
client = SparkConnectClient(chan)
self.assertEqual(client._session_id, chan.session_id)
+ def test_custom_operation_id(self):
+ client = SparkConnectClient("sc://foo/;token=bar", use_reattachable_execute=False)
+ mock = MockService(client._session_id)
+ client._stub = mock
+ req = client._execute_plan_request_with_metadata(
+ operation_id="10a4c38e-7e87-40ee-9d6f-60ff0751e63b"
+ )
+ for resp in client._stub.ExecutePlan(req, metadata=None):
+ assert resp.operation_id == "10a4c38e-7e87-40ee-9d6f-60ff0751e63b"
+
@unittest.skipIf(not should_test_connect, connect_requirement_message)
class SparkConnectClientReattachTestCase(unittest.TestCase):
diff --git a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientSuite.scala b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientSuite.scala
index ab07cf72f768..9c03d83efaf3 100644
--- a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientSuite.scala
+++ b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientSuite.scala
@@ -635,6 +635,26 @@ class SparkConnectClientSuite extends ConnectFunSuite with BeforeAndAfterEach {
observer.onNext(proto.AddArtifactsRequest.newBuilder().build())
observer.onCompleted()
}
+
+ test("client can set a custom operation id for ExecutePlan requests") {
+ startDummyServer(0)
+ client = SparkConnectClient
+ .builder()
+ .connectionString(s"sc://localhost:${server.getPort}")
+ .enableReattachableExecute()
+ .build()
+
+ val plan = buildPlan("select * from range(10000000)")
+ val dummyUUID = "10a4c38e-7e87-40ee-9d6f-60ff0751e63b"
+ val iter = client.execute(plan, operationId = Some(dummyUUID))
+ val reattachableIter =
+ ExecutePlanResponseReattachableIterator.fromIterator(iter)
+ assert(reattachableIter.operationId == dummyUUID)
+ while (reattachableIter.hasNext) {
+ val resp = reattachableIter.next()
+ assert(resp.getOperationId == dummyUUID)
+ }
+ }
}
class DummySparkConnectService() extends SparkConnectServiceGrpc.SparkConnectServiceImplBase {
diff --git a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala
index 9f8a71dd9234..e844237a3bb4 100644
--- a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala
+++ b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala
@@ -110,6 +110,15 @@ private[sql] class SparkConnectClient(
bstub.analyzePlan(request)
}
+ private def isValidUUID(uuid: String): Boolean = {
+ try {
+ UUID.fromString(uuid)
+ true
+ } catch {
+ case _: IllegalArgumentException => false
+ }
+ }
+
/**
* Execute the plan and return response iterator.
*
@@ -117,7 +126,9 @@ private[sql] class SparkConnectClient(
* done. If you don't close it, it and the underlying data will be cleaned up once the iterator
* is garbage collected.
*/
- def execute(plan: proto.Plan): CloseableIterator[proto.ExecutePlanResponse] = {
+ def execute(
+ plan: proto.Plan,
+ operationId: Option[String] = None): CloseableIterator[proto.ExecutePlanResponse] = {
artifactManager.uploadAllClassFileArtifacts()
val request = proto.ExecutePlanRequest
.newBuilder()
@@ -127,6 +138,13 @@ private[sql] class SparkConnectClient(
.setClientType(userAgent)
.addAllTags(tags.get.toSeq.asJava)
serverSideSessionId.foreach(session => request.setClientObservedServerSideSessionId(session))
+ operationId.foreach { opId =>
+ require(
+ isValidUUID(opId),
+ s"Invalid operationId: $opId. The id must be an UUID string of " +
+ "the format `00112233-4455-6677-8899-aabbccddeeff`")
+ request.setOperationId(opId)
+ }
if (configuration.useReattachableExecute) {
bstub.executePlanReattachable(request.build())
} else {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org