[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 <[email protected]>
Signed-off-by: Dongjoon Hyun <[email protected]>

Author: vicennial

python/pyspark/sql/connect/client/core.py

@@ -1228,7 +1228,9 @@ def token(self) -> Optional[str]:
         """
         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 @@ def _execute_plan_request_with_metadata(self) -> pb2.ExecutePlanRequest:
             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:

python/pyspark/sql/tests/connect/client/test_client.py

@@ -137,6 +137,7 @@ def ExecutePlan(self, req: proto.ExecutePlanRequest, metadata):
             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 @@ def test_channel_builder_with_session(self):
         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):

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 {

sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala

@@ -110,14 +110,25 @@ 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.
    *
    * It returns CloseableIterator. For resource management it is better to close it once you are
    * 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 {