diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala index 8cb5eef770a07..3555491370779 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 ``, but this was diff --git a/python/pyspark/errors/exceptions/connect.py b/python/pyspark/errors/exceptions/connect.py index 90537c2cc364e..abe436eac80c8 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 43a22c4998f2b..5c1dc029f8262 100644 --- a/python/pyspark/sql/connect/client/core.py +++ b/python/pyspark/sql/connect/client/core.py @@ -1239,7 +1239,7 @@ def to_table( 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 @@ def to_pandas( 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 @@ def execute_command( 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 @@ def _execute_plan_request_with_metadata( ) ) ) - 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 @@ def _execute_plan_request_with_metadata( 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 @@ def _execute(self, req: pb2.ExecutePlanRequest) -> None: """ 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 @@ def handle_response(b: pb2.ExecutePlanResponse) -> None: 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 @@ def _execute_and_fetch_as_iterator( # 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 @@ def handle_response( 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 @@ def clear_user_context_extensions(self) -> None: 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 @@ def _handle_error(self, error: Exception) -> NoReturn: 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 a258a4f1db70e..7af23865fed15 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 @@ def metrics(self) -> Optional[CollectedMetrics]: @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 b5bf76d86df48..4fc04f5124c3e 100644 --- a/python/pyspark/sql/tests/connect/client/test_client.py +++ b/python/pyspark/sql/tests/connect/client/test_client.py @@ -479,6 +479,33 @@ def on_execute_plan(self, req): session.client.close() session.stop() + def test_session_hook_preserves_operation_id(self): + 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 + + session = ( + RemoteSparkSession.builder.remote("sc://foo")._registerHook(TestHook).getOrCreate() + ) + try: + mock = MockService(session.client._session_id) + session.client._stub = mock + session.client.disable_reattachable_execute() + + df = session.range(1) + df.collect() + self.assertIsNotNone(df.executionInfo) + self.assertEqual(mock.req.operation_id, df.executionInfo.operation_id) + uuid.UUID(mock.req.operation_id) + finally: + session.stop() + def test_new_session_preserves_custom_channel_builder(self): class CustomChannelBuilder(DefaultChannelBuilder): pass @@ -509,6 +536,14 @@ def test_custom_operation_id(self): 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 3b5a816b7d24f..efe044a55c6ed 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 @@ def handler(**kwargs): 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 9d48eb0343d9a..49ae6ed8ea725 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 a4406c2a68fda..4a1e1c838aa84 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 f4da00b5694f0..3b6c16f06e4da 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 84b4a9000297e..8d57d35a17fda 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 9f606b698d30c..ff447ea170162 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 a433534b7511a..4735aa368587e 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.