This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch branch-4.3
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.3 by this push:
new 800438e80f36 [SPARK-58547][CONNECT] Expose operation IDs for
end-to-end request attribution
800438e80f36 is described below
commit 800438e80f3649ab3c4c84ba0b21b4a8bd3f52d1
Author: Wenchen Fan <[email protected]>
AuthorDate: Wed Aug 5 17:14:12 2026 +0800
[SPARK-58547][CONNECT] Expose operation IDs for end-to-end request
attribution
### What changes were proposed in this pull request?
Spark Connect's protocol includes an operation ID for identifying an
`ExecutePlan` request, but
clients currently leave it unset unless callers explicitly provide one.
This PR makes operation
IDs available throughout the request lifecycle:
* Python and Scala clients generate an operation ID before sending every
`ExecutePlan` request.
* Python exposes it through `ExecutionInfo.operation_id` after successful
execution and
`SparkConnectException.operation_id` after failure.
* Scala clients continue exposing it in `ExecutePlanResponse` and can
retrieve it from failures
through `SparkConnectClient.getOperationId`.
* Python session hooks preserve the client-generated operation ID.
* The server publishes the ID through
`SparkContext.SPARK_CONNECT_OPERATION_ID_PROPERTY`.
No protocol change is required, and callers that explicitly provide an
operation ID retain the
existing behavior.
### Why are the changes needed?
Successful requests can sometimes be correlated using the operation ID
returned by the server.
That is insufficient for failures before the first response, because the
caller never learns a
server-generated identifier.
Generating the identifier client-side provides a stable correlation key
before the RPC begins.
The server propagates the same ID using Spark's existing local-property
mechanism, making it
available to driver-side listeners and executor tasks. Applications and
observability integrations
can therefore correlate client failures, server logs, Spark jobs, and tasks
without depending on a
specific event system.
### Does this PR introduce _any_ user-facing change?
Yes. Spark Connect clients now generate operation IDs by default. Python
users can access the ID
through execution information and exceptions; Scala users can access it
through responses and
exceptions. Server-side integrations can access it through the public Spark
local-property key.
### How was this patch tested?
Added tests for generated operation IDs, Python session-hook preservation,
successful execution
metadata, exceptions, and server-side local-property propagation.
Ran:
* `build/sbt 'connect-client-jvm/testOnly
org.apache.spark.sql.connect.client.SparkConnectClientSuite'`
* `build/sbt 'connect/testOnly
org.apache.spark.sql.connect.service.SparkConnectServiceE2ESuite'`
The focused Python test command was also invoked, but the configured local
interpreter skipped the
Connect tests because pandas 2.2 or newer was unavailable.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex (GPT-5)
Closes #57747 from cloud-fan/cloud-fan/connect-operation-id-observability.
Authored-by: Wenchen Fan <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
(cherry picked from commit e4043f3293c6126c7d3c7ac9cb89efc976c27502)
Signed-off-by: Wenchen Fan <[email protected]>
---
.../main/scala/org/apache/spark/SparkContext.scala | 1 +
python/pyspark/errors/exceptions/connect.py | 8 ++++
python/pyspark/sql/connect/client/core.py | 33 ++++++++++------
python/pyspark/sql/metrics.py | 14 ++++++-
.../sql/tests/connect/client/test_client.py | 45 ++++++++++++++++++++++
.../sql/tests/connect/test_connect_session.py | 14 +++++++
.../connect/client/SparkConnectClientSuite.scala | 40 +++++++++++++++++++
.../client/CustomSparkConnectBlockingStub.scala | 12 ++++--
.../connect/client/GrpcExceptionConverter.scala | 19 ++++++---
.../sql/connect/client/SparkConnectClient.scala | 34 ++++++++++++----
.../connect/execution/ExecuteThreadRunner.scala | 5 ++-
.../service/SparkConnectServiceE2ESuite.scala | 31 ++++++++++++++-
12 files changed, 225 insertions(+), 31 deletions(-)
diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala
b/core/src/main/scala/org/apache/spark/SparkContext.scala
index 8cb5eef770a0..355549137077 100644
--- a/core/src/main/scala/org/apache/spark/SparkContext.scala
+++ b/core/src/main/scala/org/apache/spark/SparkContext.scala
@@ -3157,6 +3157,7 @@ object SparkContext extends Logging {
private[spark] val SQL_EXECUTION_ID_KEY = "spark.sql.execution.id"
private[spark] val DATASET_QUERY_EXECUTION_ID_KEY =
"spark.sql.dataset.queryExecution.id"
+ private[spark] val SPARK_CONNECT_OPERATION_ID_PROPERTY =
"spark.connect.operation_id"
/**
* Executor id for the driver. In earlier versions of Spark, this was
`<driver>`, but this was
diff --git a/python/pyspark/errors/exceptions/connect.py
b/python/pyspark/errors/exceptions/connect.py
index 90537c2cc364..abe436eac80c 100644
--- a/python/pyspark/errors/exceptions/connect.py
+++ b/python/pyspark/errors/exceptions/connect.py
@@ -53,6 +53,14 @@ class SparkConnectException(PySparkException):
Exception thrown from Spark Connect.
"""
+ @property
+ def operation_id(self) -> Optional[str]:
+ """The Spark Connect ExecutePlan operation ID, when available.
+
+ .. versionadded:: 4.3.0
+ """
+ return getattr(self, "_operation_id", None)
+
def convert_exception(
info: "ErrorInfo",
diff --git a/python/pyspark/sql/connect/client/core.py
b/python/pyspark/sql/connect/client/core.py
index 43a22c4998f2..5c1dc029f826 100644
--- a/python/pyspark/sql/connect/client/core.py
+++ b/python/pyspark/sql/connect/client/core.py
@@ -1239,7 +1239,7 @@ class SparkConnectClient(object):
table, schema, metrics, observed_metrics, _ =
self._execute_and_fetch(req, observations)
# Create a query execution object.
- ei = ExecutionInfo(metrics, observed_metrics)
+ ei = ExecutionInfo(metrics, observed_metrics, req.operation_id)
assert table is not None
return table, schema, ei
@@ -1275,7 +1275,7 @@ class SparkConnectClient(object):
req, observations, selfDestruct == "true"
)
assert table is not None
- ei = ExecutionInfo(metrics, observed_metrics)
+ ei = ExecutionInfo(metrics, observed_metrics, req.operation_id)
schema = schema or from_arrow_schema(table.schema,
prefer_timestamp_ntz=True)
assert schema is not None and isinstance(schema, StructType)
@@ -1418,7 +1418,7 @@ class SparkConnectClient(object):
req, observations or {}
)
# Create a query execution object.
- ei = ExecutionInfo(metrics, observed_metrics)
+ ei = ExecutionInfo(metrics, observed_metrics, req.operation_id)
if data is not None:
return (data.to_pandas(), properties, ei)
else:
@@ -1535,7 +1535,9 @@ class SparkConnectClient(object):
)
)
)
- if operation_id is not None:
+ if operation_id is None:
+ operation_id = str(uuid.uuid4())
+ else:
try:
uuid.UUID(operation_id, version=4)
except ValueError as ve:
@@ -1543,7 +1545,7 @@ class SparkConnectClient(object):
errorClass="INVALID_OPERATION_UUID_ID",
messageParameters={"arg_name": "operation_id", "origin":
str(ve)},
)
- req.operation_id = operation_id
+ req.operation_id = operation_id
self._update_request_with_user_context_extensions(req)
if call_stack_trace := self.__class__._build_call_stack_trace():
@@ -1673,8 +1675,10 @@ class SparkConnectClient(object):
"""
logger.debug("Execute")
+ operation_id = req.operation_id
for hook in self._session_hooks:
req = hook.on_execute_plan(req)
+ req.operation_id = operation_id
def handle_response(b: pb2.ExecutePlanResponse) -> None:
self._verify_response_integrity(b)
@@ -1703,7 +1707,7 @@ class SparkConnectClient(object):
for b in self._stub.ExecutePlan(req,
metadata=self._builder.metadata()):
handle_response(b)
except Exception as error:
- self._handle_error(error)
+ self._handle_error(error, req.operation_id)
def _execute_and_fetch_as_iterator(
self,
@@ -1724,8 +1728,10 @@ class SparkConnectClient(object):
# when not at debug log level.
logger.debug(f"ExecuteAndFetchAsIterator. Request:
{self._proto_to_string(req)}")
+ operation_id = req.operation_id
for hook in self._session_hooks:
req = hook.on_execute_plan(req)
+ req.operation_id = operation_id
num_records = 0
arrow_batch_chunks_to_assemble: List[bytes] = []
@@ -1932,7 +1938,7 @@ class SparkConnectClient(object):
self.interrupt_operation(req.operation_id)
raise kb
except Exception as error:
- self._handle_error(error)
+ self._handle_error(error, req.operation_id)
def _execute_and_fetch(
self,
@@ -2297,7 +2303,7 @@ class SparkConnectClient(object):
with self.global_user_context_extensions_lock:
self.global_user_context_extensions = list()
- def _handle_error(self, error: Exception) -> NoReturn:
+ def _handle_error(self, error: Exception, operation_id: Optional[str] =
None) -> NoReturn:
"""
Handle errors that occur during RPC calls.
@@ -2318,9 +2324,14 @@ class SparkConnectClient(object):
try:
self.thread_local.inside_error_handling = True
- if isinstance(error, grpc.RpcError):
- self._handle_rpc_error(error)
- raise error
+ try:
+ if isinstance(error, grpc.RpcError):
+ self._handle_rpc_error(error)
+ raise error
+ except BaseException as handled_error:
+ if operation_id:
+ handled_error._operation_id = operation_id # type:
ignore[attr-defined]
+ raise
finally:
self.thread_local.inside_error_handling = False
diff --git a/python/pyspark/sql/metrics.py b/python/pyspark/sql/metrics.py
index a258a4f1db70..7af23865fed1 100644
--- a/python/pyspark/sql/metrics.py
+++ b/python/pyspark/sql/metrics.py
@@ -297,10 +297,14 @@ class ExecutionInfo:
data frame. This value is only set in the data frame if it was executed."""
def __init__(
- self, metrics: Optional[list[PlanMetrics]], obs:
Optional[Sequence[ObservedMetrics]]
+ self,
+ metrics: Optional[list[PlanMetrics]],
+ obs: Optional[Sequence[ObservedMetrics]],
+ operation_id: Optional[str] = None,
):
self._metrics = CollectedMetrics(metrics) if metrics else None
self._observations = obs if obs else []
+ self._operation_id = operation_id
@property
def metrics(self) -> Optional[CollectedMetrics]:
@@ -309,3 +313,11 @@ class ExecutionInfo:
@property
def flows(self) -> List[Tuple[str, Dict[str, Any]]]:
return [(f.name, f.pairs) for f in self._observations]
+
+ @property
+ def operation_id(self) -> Optional[str]:
+ """The Spark Connect ExecutePlan operation ID, when available.
+
+ .. versionadded:: 4.3.0
+ """
+ return self._operation_id
diff --git a/python/pyspark/sql/tests/connect/client/test_client.py
b/python/pyspark/sql/tests/connect/client/test_client.py
index b5bf76d86df4..35252ed799ea 100644
--- a/python/pyspark/sql/tests/connect/client/test_client.py
+++ b/python/pyspark/sql/tests/connect/client/test_client.py
@@ -479,6 +479,43 @@ class SparkConnectClientTestCase(unittest.TestCase):
session.client.close()
session.stop()
+ def test_session_hook_preserves_operation_id(self):
+ execute_plan_req = None
+
+ class TestHook(RemoteSparkSession.Hook):
+ def __init__(self, _session):
+ pass
+
+ def on_execute_plan(self, req):
+ replacement = proto.ExecutePlanRequest()
+ replacement.CopyFrom(req)
+ replacement.ClearField("operation_id")
+ return replacement
+
+ class TestService(MockService):
+ def ExecutePlan(self, req, metadata, timeout=None):
+ nonlocal execute_plan_req
+ execute_plan_req = req
+ return super().ExecutePlan(req, metadata, timeout)
+
+ session = (
+
RemoteSparkSession.builder.remote("sc://foo")._registerHook(TestHook).getOrCreate()
+ )
+ try:
+ mock = TestService(session.client._session_id)
+ session.client._stub = mock
+ session.client.disable_reattachable_execute()
+
+ df = session.range(1)
+ df.collect()
+ self.assertIsNotNone(df.executionInfo)
+ self.assertIsNotNone(execute_plan_req)
+ assert execute_plan_req is not None
+ self.assertEqual(execute_plan_req.operation_id,
df.executionInfo.operation_id)
+ uuid.UUID(execute_plan_req.operation_id)
+ finally:
+ session.stop()
+
def test_new_session_preserves_custom_channel_builder(self):
class CustomChannelBuilder(DefaultChannelBuilder):
pass
@@ -509,6 +546,14 @@ class SparkConnectClientTestCase(unittest.TestCase):
for resp in client._stub.ExecutePlan(req, metadata=None):
assert resp.operation_id == "10a4c38e-7e87-40ee-9d6f-60ff0751e63b"
+ def test_execute_plan_request_generates_operation_id(self):
+ client = SparkConnectClient("sc://foo/;token=bar",
use_reattachable_execute=False)
+ try:
+ req = client._execute_plan_request_with_metadata()
+ uuid.UUID(req.operation_id)
+ finally:
+ client.close()
+
def test_on_exit_calls_release_and_close_when_enabled(self):
client = SparkConnectClient("sc://foo/",
use_reattachable_execute=False)
client._release_session_on_exit = True
diff --git a/python/pyspark/sql/tests/connect/test_connect_session.py
b/python/pyspark/sql/tests/connect/test_connect_session.py
index 3b5a816b7d24..efe044a55c6e 100644
--- a/python/pyspark/sql/tests/connect/test_connect_session.py
+++ b/python/pyspark/sql/tests/connect/test_connect_session.py
@@ -82,6 +82,20 @@ class SparkConnectSessionTests(ReusedConnectTestCase):
self.spark.sql("select 1").collect()
self.assertGreaterEqual(len(handler_called), 0)
+ @timeout(10)
+ def test_operation_id_in_execution_info_and_exception(self):
+ df = self.spark.sql("select 1")
+ df.collect()
+ self.assertIsNotNone(df.executionInfo)
+ operation_id = df.executionInfo.operation_id
+ self.assertIsNotNone(operation_id)
+ uuid.UUID(operation_id)
+
+ with self.assertRaises(SparkConnectException) as error:
+ self.spark.sql("select raise_error('expected')").collect()
+ self.assertIsNotNone(error.exception.operation_id)
+ uuid.UUID(error.exception.operation_id)
+
def _check_no_active_session_error(self, e: PySparkException):
self.check_error(exception=e, errorClass="NO_ACTIVE_SESSION",
messageParameters=dict())
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 9d48eb0343d9..49ae6ed8ea72 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
@@ -78,6 +78,46 @@ class SparkConnectClientSuite extends ConnectFunSuite {
assert(client.userId == System.getProperty("user.name"))
}
+ test("client generates an operation ID for ExecutePlan requests") {
+ startDummyServer(0)
+ client = SparkConnectClient
+ .builder()
+ .connectionString(s"sc://localhost:${server.getPort}")
+ .disableReattachableExecute()
+ .build()
+
+ val responses = client.execute(buildPlan("select 1")).toSeq
+ val operationId = responses.head.getOperationId
+
+ UUID.fromString(operationId)
+ assert(responses.forall(_.getOperationId == operationId))
+ }
+
+ test("ExecutePlan exceptions expose the client-generated operation ID") {
+ val failingService = new DummySparkConnectService {
+ override def executePlan(
+ request: ExecutePlanRequest,
+ responseObserver: StreamObserver[ExecutePlanResponse]): Unit = {
+
responseObserver.onError(Status.INTERNAL.withDescription("expected").asRuntimeException())
+ }
+ }
+ server =
NettyServerBuilder.forPort(0).addService(failingService).build().start()
+ service = failingService
+ client = SparkConnectClient
+ .builder()
+ .connectionString(s"sc://localhost:${server.getPort}")
+ .disableReattachableExecute()
+ .retryPolicy(RetryPolicy(maxRetries = Some(0), canRetry = _ => false,
name = "NoRetry"))
+ .build()
+
+ val error = intercept[SparkException] {
+ client.execute(buildPlan("select 1")).foreach(_ => ())
+ }
+ val operationId = SparkConnectClient.getOperationId(error)
+ assert(operationId.isDefined)
+ UUID.fromString(operationId.get)
+ }
+
test("Placeholder test: Create SparkConnectClient") {
client = SparkConnectClient.builder().userId("abc123").build()
assert(client.userId == "abc123")
diff --git
a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/CustomSparkConnectBlockingStub.scala
b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/CustomSparkConnectBlockingStub.scala
index a4406c2a68fd..4a1e1c838aa8 100644
---
a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/CustomSparkConnectBlockingStub.scala
+++
b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/CustomSparkConnectBlockingStub.scala
@@ -52,7 +52,8 @@ private[connect] class CustomSparkConnectBlockingStub(
grpcExceptionConverter.convert(
request.getSessionId,
request.getUserContext,
- request.getClientType) {
+ request.getClientType,
+ Option(request.getOperationId).filter(_.nonEmpty)) {
grpcExceptionConverter.convertIterator[ExecutePlanResponse](
request.getSessionId,
request.getUserContext,
@@ -62,7 +63,8 @@ private[connect] class CustomSparkConnectBlockingStub(
r => {
stubState.responseValidator.wrapIterator(
CloseableIterator(stub.executePlan(r).asScala))
- }))
+ }),
+ Option(request.getOperationId).filter(_.nonEmpty))
}
}
@@ -71,7 +73,8 @@ private[connect] class CustomSparkConnectBlockingStub(
grpcExceptionConverter.convert(
request.getSessionId,
request.getUserContext,
- request.getClientType) {
+ request.getClientType,
+ Option(request.getOperationId).filter(_.nonEmpty)) {
grpcExceptionConverter.convertIterator[ExecutePlanResponse](
request.getSessionId,
request.getUserContext,
@@ -83,7 +86,8 @@ private[connect] class CustomSparkConnectBlockingStub(
channel,
stubState.retryHandler,
stubState.rpcDeadlines.reattachableExecutePlan,
- stubState.rpcDeadlines.reattachExecute)))
+ stubState.rpcDeadlines.reattachExecute)),
+ Option(request.getOperationId).filter(_.nonEmpty))
}
}
diff --git
a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/GrpcExceptionConverter.scala
b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/GrpcExceptionConverter.scala
index f4da00b5694f..3b6c16f06e4d 100644
---
a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/GrpcExceptionConverter.scala
+++
b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/GrpcExceptionConverter.scala
@@ -64,12 +64,18 @@ private[client] class GrpcExceptionConverter(
.map(d => grpcStub.withDeadline(Deadline.after(d.toMillis,
TimeUnit.MILLISECONDS)))
.getOrElse(grpcStub)
- def convert[T](sessionId: String, userContext: UserContext, clientType:
String)(f: => T): T = {
+ def convert[T](
+ sessionId: String,
+ userContext: UserContext,
+ clientType: String,
+ operationId: Option[String] = None)(f: => T): T = {
try {
f
} catch {
case e: StatusRuntimeException =>
- throw toThrowable(e, sessionId, userContext, clientType)
+ val converted = toThrowable(e, sessionId, userContext, clientType)
+ operationId.foreach(SparkConnectClient.attachOperationId(converted, _))
+ throw converted
}
}
@@ -77,25 +83,26 @@ private[client] class GrpcExceptionConverter(
sessionId: String,
userContext: UserContext,
clientType: String,
- iter: CloseableIterator[T]): CloseableIterator[T] = {
+ iter: CloseableIterator[T],
+ operationId: Option[String] = None): CloseableIterator[T] = {
new WrappedCloseableIterator[T] {
override def innerIterator: Iterator[T] = iter
override def hasNext: Boolean = {
- convert(sessionId, userContext, clientType) {
+ convert(sessionId, userContext, clientType, operationId) {
iter.hasNext
}
}
override def next(): T = {
- convert(sessionId, userContext, clientType) {
+ convert(sessionId, userContext, clientType, operationId) {
iter.next()
}
}
override def close(): Unit = {
- convert(sessionId, userContext, clientType) {
+ convert(sessionId, userContext, clientType, operationId) {
iter.close()
}
}
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 eeb14fcfc2c7..905e9862d2a7 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
@@ -318,13 +318,12 @@ private[sql] class SparkConnectClient(
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)
- }
+ val resolvedOperationId = operationId.getOrElse(UUID.randomUUID.toString)
+ require(
+ isValidUUID(resolvedOperationId),
+ s"Invalid operationId: $resolvedOperationId. The id must be an UUID
string of " +
+ "the format `00112233-4455-6677-8899-aabbccddeeff`")
+ request.setOperationId(resolvedOperationId)
if (configuration.useReattachableExecute) {
bstub.executePlanReattachable(request.build())
} else {
@@ -699,8 +698,29 @@ private[sql] class SparkConnectClient(
// Options for plan compression
case class PlanCompressionOptions(thresholdBytes: Int, algorithm: String)
+private final class SparkConnectOperationIdException(val operationId: String)
+ extends RuntimeException(s"Spark Connect operation ID: $operationId",
null, false, false)
+
object SparkConnectClient {
+ /**
+ * Returns the ExecutePlan operation ID attached to a Spark Connect failure,
when available.
+ *
+ * @since 4.3.0
+ */
+ @DeveloperApi
+ def getOperationId(error: Throwable): Option[String] = {
+ error.getSuppressed.collectFirst { case marker:
SparkConnectOperationIdException =>
+ marker.operationId
+ }
+ }
+
+ private[client] def attachOperationId(error: Throwable, operationId:
String): Unit = {
+ if (getOperationId(error).isEmpty) {
+ error.addSuppressed(new SparkConnectOperationIdException(operationId))
+ }
+ }
+
private[sql] val SPARK_REMOTE: String = "SPARK_REMOTE"
private val DEFAULT_USER_AGENT: String = "_SPARK_CONNECT_SCALA"
diff --git
a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/ExecuteThreadRunner.scala
b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/ExecuteThreadRunner.scala
index 9f606b698d30..ff447ea17016 100644
---
a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/ExecuteThreadRunner.scala
+++
b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/ExecuteThreadRunner.scala
@@ -25,7 +25,7 @@ import scala.util.control.NonFatal
import com.google.protobuf.Message
-import org.apache.spark.SparkSQLException
+import org.apache.spark.{SparkContext, SparkSQLException}
import org.apache.spark.connect.proto
import org.apache.spark.internal.{Logging, LogKeys}
import org.apache.spark.sql.connect.common.ProtoUtils
@@ -224,6 +224,9 @@ private[connect] class ExecuteThreadRunner(executeHolder:
ExecuteHolder) extends
"callSite.short",
s"Spark Connect - ${Utils.abbreviate(debugString, 128)}")
session.sparkContext.setLocalProperty("callSite.long",
Utils.abbreviate(debugString, 2048))
+ session.sparkContext.setLocalProperty(
+ SparkContext.SPARK_CONNECT_OPERATION_ID_PROPERTY,
+ executeHolder.operationId)
executeHolder.request.getPlan.getOpTypeCase match {
case proto.Plan.OpTypeCase.ROOT | proto.Plan.OpTypeCase.COMMAND =>
diff --git
a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceE2ESuite.scala
b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceE2ESuite.scala
index a433534b7511..56e69486e1fa 100644
---
a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceE2ESuite.scala
+++
b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceE2ESuite.scala
@@ -18,14 +18,17 @@ package org.apache.spark.sql.connect.service
import java.io.ByteArrayOutputStream
import java.util.UUID
+import java.util.concurrent.{CountDownLatch, TimeUnit}
+import java.util.concurrent.atomic.AtomicReference
import com.github.luben.zstd.{Zstd, ZstdOutputStreamNoFinalizer}
import com.google.protobuf.ByteString
import org.scalatest.concurrent.Eventually
import org.scalatest.time.SpanSugar._
-import org.apache.spark.SparkException
+import org.apache.spark.{SparkContext, SparkException}
import org.apache.spark.connect.proto
+import org.apache.spark.scheduler.{SparkListener, SparkListenerJobStart}
import org.apache.spark.sql.connect.SparkConnectServerTest
import org.apache.spark.sql.connect.config.Connect
@@ -37,6 +40,32 @@ class SparkConnectServiceE2ESuite extends
SparkConnectServerTest {
// were all already in the buffer.
val BIG_ENOUGH_QUERY = "select * from range(1000000)"
+ test("ExecutePlan operation ID is available as a Spark local property") {
+ val operationIdFromJob = new AtomicReference[String]()
+ val jobStarted = new CountDownLatch(1)
+ val listener = new SparkListener {
+ override def onJobStart(jobStart: SparkListenerJobStart): Unit = {
+ val operationId =
+
jobStart.properties.getProperty(SparkContext.SPARK_CONNECT_OPERATION_ID_PROPERTY)
+ if (operationId != null) {
+ operationIdFromJob.set(operationId)
+ jobStarted.countDown()
+ }
+ }
+ }
+ spark.sparkContext.addSparkListener(listener)
+ try {
+ withClient { client =>
+ val responses = client.execute(buildPlan("select count(*) from
range(10)")).toSeq
+ val operationId = responses.head.getOperationId
+ assert(jobStarted.await(10, TimeUnit.SECONDS))
+ assert(operationIdFromJob.get() == operationId)
+ }
+ } finally {
+ spark.sparkContext.removeSparkListener(listener)
+ }
+ }
+
test("Execute is sent eagerly to the server upon iterator creation") {
// This behavior changed with grpc upgrade from 1.56.0 to 1.59.0.
// Testing to be aware of future changes.
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
