You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ru...@apache.org on 2026/09/11 13:35:37 UTC
(spark) branch master updated: [SPARK-59135][CONNECT][PYTHON] Handle metadata in command response streams
This is an automated email from the ASF dual-hosted git repository.
zhengruifeng 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 29a4be87d653 [SPARK-59135][CONNECT][PYTHON] Handle metadata in command response streams
29a4be87d653 is described below
commit 29a4be87d653da8efd7e0242ce3ccf1358b938dd
Author: haoyangeng-db <ha...@gmail.com>
AuthorDate: Fri Sep 11 21:35:24 2026 +0800
[SPARK-59135][CONNECT][PYTHON] Handle metadata in command response streams
### What changes were proposed in this pull request?
This PR makes `SparkConnectClient.execute_command_as_iterator` yield every decoded response item instead of requiring every item to be a dictionary. Streaming-query listener and pipeline consumers ignore `PlanMetrics` and `ObservedMetrics` while continuing to reject unrelated malformed responses.
It also creates Python Data Source profiler accumulators only when performance or memory profiling is enabled.
### Why are the changes needed?
Python Data Source workers currently register profiler accumulators even when profiling is disabled. Their empty accumulator updates can arrive alongside command results as observed metrics. The Spark Connect command iterator treats those valid non-dictionary responses as unknown and terminates long-lived command streams, including streaming listeners.
The low-level iterator already decodes several legitimate response types, so it should preserve them and let each command-specific consumer decide which types it accepts.
### Does this PR introduce _any_ user-facing change?
Yes. Spark Connect command streams no longer fail when valid execution metrics accompany command results. Python Data Source workers also stop sending empty profiler updates when profiling is disabled. Unexpected response shapes still raise an error.
### How was this patch tested?
Added focused tests covering:
- command results followed by plan and observed metrics;
- protobuf extension responses;
- streaming-listener metadata before and after registration;
- malformed listener and pipeline responses;
- disabled, unsupported, performance, and memory profiler modes;
- preservation of ordinary accumulator updates.
All 7 focused tests and all 44 tests in the mocked Spark Connect client suite passed. Ruff check/format and `git diff --check` passed.
A full SBT build could not resolve external dependencies in this environment because DNS was unavailable.
### Was this patch authored or co-authored using generative AI tooling?
Co-authored with Codex
Closes #58435 from haoyangeng-db/fix-connect-command-iterator-metadata.
Authored-by: haoyangeng-db <ha...@gmail.com>
Signed-off-by: Ruifeng Zheng <ru...@foxmail.com>
---
dev/sparktestsupport/modules.py | 1 +
python/pyspark/pipelines/spark_connect_pipeline.py | 20 +++--
.../pyspark/pipelines/tests/test_spark_connect.py | 35 +++++++++
python/pyspark/sql/connect/client/core.py | 50 +++++--------
python/pyspark/sql/connect/streaming/query.py | 41 +++++++---
.../sql/tests/connect/client/test_client.py | 50 ++++++++++++-
.../sql/tests/connect/streaming/test_listener.py | 80 ++++++++++++++++++++
python/pyspark/sql/tests/test_python_datasource.py | 87 ++++++++++++++++++++++
python/pyspark/sql/worker/utils.py | 49 ++++++------
9 files changed, 339 insertions(+), 74 deletions(-)
diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py
index d214f8a514b9..8441926d2cf3 100644
--- a/dev/sparktestsupport/modules.py
+++ b/dev/sparktestsupport/modules.py
@@ -1368,6 +1368,7 @@ pyspark_structured_streaming_connect = Module(
python_test_goals=[
# unittests
"pyspark.sql.tests.connect.test_parity_python_streaming_datasource",
+ "pyspark.sql.tests.connect.streaming.test_listener",
"pyspark.sql.tests.connect.streaming.test_parity_streaming",
"pyspark.sql.tests.connect.streaming.test_parity_listener",
"pyspark.sql.tests.connect.streaming.test_parity_foreach",
diff --git a/python/pyspark/pipelines/spark_connect_pipeline.py b/python/pyspark/pipelines/spark_connect_pipeline.py
index 6b5858c08729..91e5e70efc8c 100644
--- a/python/pyspark/pipelines/spark_connect_pipeline.py
+++ b/python/pyspark/pipelines/spark_connect_pipeline.py
@@ -15,12 +15,16 @@
# limitations under the License.
#
from datetime import timezone
-from typing import Any, Dict, Iterator, Mapping, Optional, Sequence, cast
+from typing import TYPE_CHECKING, Any, Iterator, Mapping, Optional, Sequence, cast
import pyspark.sql.connect.proto as pb2
from pyspark.errors.exceptions.base import PySparkValueError
from pyspark.pipelines.logging_utils import log_with_provided_timestamp
from pyspark.sql import SparkSession
+from pyspark.sql.metrics import ObservedMetrics, PlanMetrics
+
+if TYPE_CHECKING:
+ from pyspark.sql.connect.client.core import _ExecutePlanResponseItem
def create_dataflow_graph(
@@ -46,16 +50,22 @@ def create_dataflow_graph(
return properties["pipeline_command_result"].create_dataflow_graph_result.dataflow_graph_id
-def handle_pipeline_events(iter: Iterator[Dict[str, Any]]) -> None:
+def handle_pipeline_events(iter: Iterator["_ExecutePlanResponseItem"]) -> None:
"""
Prints out the pipeline events received from the Spark Connect server.
"""
for result in iter:
- if "pipeline_command_result" in result.keys():
+ if isinstance(result, (PlanMetrics, ObservedMetrics)):
+ continue
+ elif not isinstance(result, dict):
+ raise PySparkValueError(
+ f"Pipeline logs stream handler received an unexpected result: {result}"
+ )
+ elif "pipeline_command_result" in result:
# We expect to get a pipeline_command_result back in response to the initial StartRun
# command.
continue
- elif "pipeline_event_result" not in result.keys():
+ elif "pipeline_event_result" not in result:
raise PySparkValueError(
f"Pipeline logs stream handler received an unexpected result: {result}"
)
@@ -73,7 +83,7 @@ def start_run(
refresh: Optional[Sequence[str]],
dry: bool,
storage: str,
-) -> Iterator[Dict[str, Any]]:
+) -> Iterator["_ExecutePlanResponseItem"]:
"""Start a run of the dataflow graph in the Spark Connect server.
:param spark: SparkSession.
diff --git a/python/pyspark/pipelines/tests/test_spark_connect.py b/python/pyspark/pipelines/tests/test_spark_connect.py
index 9cd0376b0549..8117ec6d28e1 100644
--- a/python/pyspark/pipelines/tests/test_spark_connect.py
+++ b/python/pyspark/pipelines/tests/test_spark_connect.py
@@ -20,6 +20,7 @@ Tests that run Pipelines against a Spark Connect server.
"""
import unittest
+from unittest import mock
from pyspark import pipelines as dp
from pyspark.testing.connectutils import (
@@ -29,6 +30,10 @@ from pyspark.testing.connectutils import (
)
if should_test_connect:
+ from google.protobuf import any_pb2
+
+ import pyspark.sql.connect.proto as pb2
+ from pyspark.errors import PySparkValueError
from pyspark.errors.exceptions.connect import AnalysisException
from pyspark.pipelines.graph_element_registry import graph_element_registration_context
from pyspark.pipelines.spark_connect_graph_element_registry import (
@@ -39,6 +44,36 @@ if should_test_connect:
handle_pipeline_events,
start_run,
)
+ from pyspark.sql.connect.client.core import PlanObservedMetrics
+ from pyspark.sql.metrics import PlanMetrics
+
+
+@unittest.skipIf(not should_test_connect, connect_requirement_message)
+class SparkConnectPipelineEventHandlerTests(unittest.TestCase):
+ def test_ignores_execution_metadata(self):
+ pipeline_event = pb2.PipelineEventResult()
+ pipeline_event.event.message = "event"
+ results = iter(
+ [
+ {"pipeline_command_result": object()},
+ {"pipeline_event_result": pipeline_event},
+ PlanMetrics("metric", 1, 0, []),
+ PlanObservedMetrics("observation", [], []),
+ ]
+ )
+
+ with mock.patch(
+ "pyspark.pipelines.spark_connect_pipeline.log_with_provided_timestamp"
+ ) as log:
+ handle_pipeline_events(results)
+
+ log.assert_called_once_with("event", mock.ANY)
+
+ def test_rejects_unexpected_response(self):
+ with self.assertRaises(PySparkValueError) as error:
+ handle_pipeline_events(iter([any_pb2.Any()]))
+
+ self.assertIn("received an unexpected result", str(error.exception))
@unittest.skipIf(not should_test_connect, connect_requirement_message)
diff --git a/python/pyspark/sql/connect/client/core.py b/python/pyspark/sql/connect/client/core.py
index e47bc3c1ae10..e5826887f195 100644
--- a/python/pyspark/sql/connect/client/core.py
+++ b/python/pyspark/sql/connect/client/core.py
@@ -694,6 +694,16 @@ class PlanObservedMetrics(ObservedMetrics):
}
+_ExecutePlanResponseItem = Union[
+ "pa.RecordBatch",
+ StructType,
+ PlanMetrics,
+ PlanObservedMetrics,
+ Dict[str, Any],
+ any_pb2.Any,
+]
+
+
class AnalyzeResult:
def __init__(
self,
@@ -1447,9 +1457,11 @@ class SparkConnectClient(object):
def execute_command_as_iterator(
self, command: pb2.Command, observations: Optional[Dict[str, Observation]] = None
- ) -> Iterator[Dict[str, Any]]:
+ ) -> Iterator[_ExecutePlanResponseItem]:
"""
- Execute given command. Similar to execute_command, but the value is returned using yield.
+ Execute given command, yielding each decoded response as it arrives.
+
+ Callers are responsible for handling the response types relevant to their command.
"""
if logger.isEnabledFor(logging.DEBUG):
# inside an if statement to not incur a performance cost converting proto to string
@@ -1459,16 +1471,7 @@ class SparkConnectClient(object):
)
req = self._execute_plan_request_with_metadata()
self._set_command_in_plan(req.plan, command)
- for response in self._execute_and_fetch_as_iterator(req, observations or {}):
- if isinstance(response, dict):
- yield response
- else:
- raise PySparkValueError(
- errorClass="UNKNOWN_RESPONSE",
- messageParameters={
- "response": str(response),
- },
- )
+ yield from self._execute_and_fetch_as_iterator(req, observations or {})
def same_semantics(self, plan: pb2.Plan, other: pb2.Plan) -> bool:
"""
@@ -1749,15 +1752,7 @@ class SparkConnectClient(object):
req: pb2.ExecutePlanRequest,
observations: Dict[str, Observation],
progress: Optional["Progress"] = None,
- ) -> Iterator[
- Union[
- "pa.RecordBatch",
- StructType,
- PlanMetrics,
- PlanObservedMetrics,
- Dict[str, Any],
- ]
- ]:
+ ) -> Iterator[_ExecutePlanResponseItem]:
if logger.isEnabledFor(logging.DEBUG):
# inside an if statement to not incur a performance cost converting proto to string
# when not at debug log level.
@@ -1772,16 +1767,7 @@ class SparkConnectClient(object):
def handle_response(
b: pb2.ExecutePlanResponse,
- ) -> Iterator[
- Union[
- "pa.RecordBatch",
- StructType,
- PlanMetrics,
- PlanObservedMetrics,
- Dict[str, Any],
- any_pb2.Any,
- ]
- ]:
+ ) -> Iterator[_ExecutePlanResponseItem]:
nonlocal num_records
# The session ID is the local session ID and should match what we expect.
self._verify_response_integrity(b)
@@ -2014,7 +2000,7 @@ class SparkConnectClient(object):
raise PySparkValueError(
errorClass="UNKNOWN_RESPONSE",
messageParameters={
- "response": response,
+ "response": str(response),
},
)
diff --git a/python/pyspark/sql/connect/streaming/query.py b/python/pyspark/sql/connect/streaming/query.py
index e5d16073f4d5..ea1a7b07bfd6 100644
--- a/python/pyspark/sql/connect/streaming/query.py
+++ b/python/pyspark/sql/connect/streaming/query.py
@@ -26,6 +26,7 @@ from pyspark.errors.exceptions.connect import (
StreamingQueryException as CapturedStreamingQueryException,
)
from pyspark.sql.connect import proto
+from pyspark.sql.metrics import ObservedMetrics, PlanMetrics
from pyspark.sql.streaming import StreamingQueryListener
from pyspark.sql.streaming.listener import (
QueryIdleEvent,
@@ -42,6 +43,7 @@ from pyspark.sql.streaming.query import (
)
if TYPE_CHECKING:
+ from pyspark.sql.connect.client.core import _ExecutePlanResponseItem
from pyspark.sql.connect.session import SparkSession
@@ -342,7 +344,30 @@ class StreamingQueryListenerBus:
self._listener_bus.remove(listener)
- def _register_server_side_listener(self) -> Iterator[Dict[str, Any]]:
+ @staticmethod
+ def _iter_listener_events(
+ results: Iterator["_ExecutePlanResponseItem"],
+ ) -> Iterator[pb2.StreamingQueryListenerEventsResult]:
+ for result in results:
+ if isinstance(result, (PlanMetrics, ObservedMetrics)):
+ continue
+ if (
+ not isinstance(result, dict)
+ or "streaming_query_listener_events_result" not in result
+ ):
+ raise PySparkValueError(
+ errorClass="UNKNOWN_RESPONSE",
+ messageParameters={"response": str(result)},
+ )
+ response = result["streaming_query_listener_events_result"]
+ if not isinstance(response, pb2.StreamingQueryListenerEventsResult):
+ raise PySparkValueError(
+ errorClass="UNKNOWN_RESPONSE",
+ messageParameters={"response": str(result)},
+ )
+ yield response
+
+ def _register_server_side_listener(self) -> Iterator["_ExecutePlanResponseItem"]:
"""
Send add listener request to the server, after received confirmation from the server,
start a new thread to handle these events.
@@ -353,27 +378,19 @@ class StreamingQueryListenerBus:
exec_cmd.streaming_query_listener_bus_command.CopyFrom(cmd)
result_iter = self._sqm._session.client.execute_command_as_iterator(exec_cmd)
# Main thread should block until received listener_added_success message
- for result in result_iter:
- response = cast(
- pb2.StreamingQueryListenerEventsResult,
- result["streaming_query_listener_events_result"],
- )
+ for response in self._iter_listener_events(result_iter):
if response.HasField("listener_bus_listener_added"):
break
return result_iter
- def _query_event_handler(self, iter: Iterator[Dict[str, Any]]) -> None:
+ def _query_event_handler(self, iter: Iterator["_ExecutePlanResponseItem"]) -> None:
"""
Handler function passed to the new thread, if there is any error while receiving
listener events, it means the connection is unstable. In this case, remove all listeners
and tell the user to add back the listeners.
"""
try:
- for result in iter:
- response = cast(
- pb2.StreamingQueryListenerEventsResult,
- result["streaming_query_listener_events_result"],
- )
+ for response in self._iter_listener_events(iter):
for event in response.events:
deserialized_event = self.deserialize(event)
self.post_to_all(deserialized_event)
diff --git a/python/pyspark/sql/tests/connect/client/test_client.py b/python/pyspark/sql/tests/connect/client/test_client.py
index b1463642dfdc..b3f32f0a78f1 100644
--- a/python/pyspark/sql/tests/connect/client/test_client.py
+++ b/python/pyspark/sql/tests/connect/client/test_client.py
@@ -36,13 +36,14 @@ if should_test_connect:
from pyspark.errors import PySparkRuntimeError
from pyspark.errors.exceptions.connect import SparkConnectGrpcException
from pyspark.sql.connect.client import DefaultChannelBuilder, SparkConnectClient
- from pyspark.sql.connect.client.core import RpcDeadlines
+ from pyspark.sql.connect.client.core import PlanObservedMetrics, RpcDeadlines
from pyspark.sql.connect.client.reattach import ExecutePlanResponseReattachableIterator
from pyspark.sql.connect.client.retries import (
DefaultPolicy,
Retrying,
)
from pyspark.sql.connect.session import SparkSession as RemoteSparkSession
+ from pyspark.sql.metrics import PlanMetrics
class TestPolicy(DefaultPolicy):
def __init__(self):
@@ -153,11 +154,12 @@ if should_test_connect:
),
]
- def __init__(self, session_id: str, operation_statuses=None):
+ def __init__(self, session_id: str, operation_statuses=None, execute_plan_responses=None):
self._session_id = session_id
self.req = None
self.metadata = None
self.client_user_context_extensions = []
+ self._execute_plan_responses = execute_plan_responses
if operation_statuses is None:
operation_statuses = self.DEFAULT_OPERATION_STATUSES
self._operation_statuses = {s.operation_id: s for s in operation_statuses}
@@ -166,6 +168,12 @@ if should_test_connect:
self.req = req
self.metadata = metadata
self.client_user_context_extensions = list(req.user_context.extensions)
+ if self._execute_plan_responses is not None:
+ for response in self._execute_plan_responses:
+ response.session_id = self._session_id
+ response.operation_id = req.operation_id
+ return self._execute_plan_responses
+
resp = proto.ExecutePlanResponse()
resp.session_id = self._session_id
resp.operation_id = req.operation_id
@@ -262,6 +270,44 @@ if should_test_connect:
@unittest.skipIf(not should_test_connect, connect_requirement_message)
class SparkConnectClientTestCase(unittest.TestCase):
+ def test_command_iterator_yields_all_responses(self):
+ command_response = proto.ExecutePlanResponse()
+ command_response.pipeline_command_result.SetInParent()
+ metadata_response = proto.ExecutePlanResponse()
+ metadata_response.metrics.metrics.add(name="metric", plan_id=1)
+ metadata_response.observed_metrics.add(name="observation")
+
+ client = SparkConnectClient("sc://foo/", use_reattachable_execute=False)
+ client._stub = MockService(
+ client._session_id,
+ execute_plan_responses=[command_response, metadata_response],
+ )
+
+ results = list(client.execute_command_as_iterator(proto.Command()))
+
+ self.assertEqual(len(results), 3)
+ self.assertEqual(list(results[0]), ["pipeline_command_result"])
+ self.assertIsInstance(results[1], PlanMetrics)
+ self.assertIsInstance(results[2], PlanObservedMetrics)
+
+ def test_command_iterator_yields_extension(self):
+ response = proto.ExecutePlanResponse()
+ response.extension.Pack(wrappers_pb2.StringValue(value="unexpected"))
+
+ client = SparkConnectClient("sc://foo/", use_reattachable_execute=False)
+ client._stub = MockService(
+ client._session_id,
+ execute_plan_responses=[response],
+ )
+
+ results = list(client.execute_command_as_iterator(proto.Command()))
+
+ self.assertEqual(len(results), 1)
+ self.assertIsInstance(results[0], any_pb2.Any)
+ value = wrappers_pb2.StringValue()
+ self.assertTrue(results[0].Unpack(value))
+ self.assertEqual(value.value, "unexpected")
+
def test_user_agent_passthrough(self):
client = SparkConnectClient("sc://foo/;user_agent=bar", use_reattachable_execute=False)
mock = MockService(client._session_id)
diff --git a/python/pyspark/sql/tests/connect/streaming/test_listener.py b/python/pyspark/sql/tests/connect/streaming/test_listener.py
new file mode 100644
index 000000000000..1c15f913e76f
--- /dev/null
+++ b/python/pyspark/sql/tests/connect/streaming/test_listener.py
@@ -0,0 +1,80 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import unittest
+from unittest import mock
+
+from pyspark.errors import PySparkValueError
+from pyspark.testing.connectutils import connect_requirement_message, should_test_connect
+
+if should_test_connect:
+ from google.protobuf import any_pb2
+
+ import pyspark.sql.connect.proto as pb2
+ from pyspark.sql.connect.client.core import PlanObservedMetrics
+ from pyspark.sql.connect.streaming.query import StreamingQueryListenerBus
+ from pyspark.sql.metrics import PlanMetrics
+
+
+@unittest.skipIf(not should_test_connect, connect_requirement_message)
+class StreamingQueryListenerBusResponseTests(unittest.TestCase):
+ def test_ignores_execution_metadata_before_and_after_registration(self):
+ listener_added = pb2.StreamingQueryListenerEventsResult(listener_bus_listener_added=True)
+ listener_events = pb2.StreamingQueryListenerEventsResult()
+ event = listener_events.events.add(event_json="{}")
+ results = iter(
+ [
+ PlanMetrics("metric", 1, 0, []),
+ {"streaming_query_listener_events_result": listener_added},
+ PlanObservedMetrics("observation", [], []),
+ {"streaming_query_listener_events_result": listener_events},
+ ]
+ )
+ manager = mock.Mock()
+ manager._session.client.execute_command_as_iterator.return_value = results
+ listener_bus = StreamingQueryListenerBus(manager)
+
+ remaining_results = listener_bus._register_server_side_listener()
+ with (
+ mock.patch.object(
+ listener_bus, "deserialize", return_value=mock.sentinel.deserialized_event
+ ) as deserialize,
+ mock.patch.object(listener_bus, "post_to_all") as post_to_all,
+ ):
+ listener_bus._query_event_handler(remaining_results)
+
+ deserialize.assert_called_once_with(event)
+ post_to_all.assert_called_once_with(mock.sentinel.deserialized_event)
+
+ def test_rejects_unexpected_response(self):
+ responses = [
+ any_pb2.Any(),
+ {"unexpected": object()},
+ {"streaming_query_listener_events_result": object()},
+ ]
+ for response in responses:
+ with self.subTest(response=response):
+ with self.assertRaises(PySparkValueError) as error:
+ list(StreamingQueryListenerBus._iter_listener_events(iter([response])))
+
+ self.assertEqual(error.exception.getCondition(), "UNKNOWN_RESPONSE")
+
+
+if __name__ == "__main__":
+ from pyspark.testing import main
+
+ main()
diff --git a/python/pyspark/sql/tests/test_python_datasource.py b/python/pyspark/sql/tests/test_python_datasource.py
index 5bb8e9df1e3b..b9e2574879ea 100644
--- a/python/pyspark/sql/tests/test_python_datasource.py
+++ b/python/pyspark/sql/tests/test_python_datasource.py
@@ -25,6 +25,7 @@ from dataclasses import dataclass
from datetime import datetime
from decimal import Decimal
from typing import Callable, Iterable, Iterator, List, Tuple, Union
+from unittest import mock
from pyspark.errors import AnalysisException, PythonException
from pyspark.memory_profiler_ext import has_memory_profiler
@@ -66,6 +67,92 @@ from pyspark.testing.utils import (
from pyspark.util import is_remote_only
+class PythonDataSourceWorkerUtilsTests(unittest.TestCase):
+ def test_profiler_accumulators_are_created_only_when_profiling_is_enabled(self):
+ from pyspark.accumulators import (
+ INT_ACCUMULATOR_PARAM,
+ SpecialAccumulatorIds,
+ _accumulatorRegistry,
+ _deserialize_accumulator,
+ )
+ from pyspark.serializers import SpecialLengths
+ from pyspark.serializers import read_int as read_serialized_int
+
+ with mock.patch.dict(os.environ, {"SPARK_PYTHON_RUNTIME": "PYTHON_WORKER"}):
+ from pyspark.sql.worker import utils as worker_utils
+
+ self.addCleanup(_accumulatorRegistry.clear)
+
+ def run_worker(profiler, regular_accumulator_id=None):
+ infile = io.BytesIO()
+ outfile = io.BytesIO()
+ main_call_count = 0
+
+ def main(_infile, _outfile):
+ nonlocal main_call_count
+ main_call_count += 1
+ if regular_accumulator_id is not None:
+ _deserialize_accumulator(
+ regular_accumulator_id,
+ 0,
+ INT_ACCUMULATOR_PARAM,
+ )
+
+ with mock.patch.multiple(
+ worker_utils,
+ check_python_version=mock.DEFAULT,
+ start_faulthandler_periodic_traceback=mock.DEFAULT,
+ setup_memory_limits=mock.DEFAULT,
+ setup_spark_files=mock.DEFAULT,
+ setup_broadcasts=mock.DEFAULT,
+ RunnerConf=mock.DEFAULT,
+ read_int=mock.DEFAULT,
+ WorkerPerfProfiler=mock.DEFAULT,
+ WorkerMemoryProfiler=mock.DEFAULT,
+ ) as patched:
+ patched["RunnerConf"].return_value.profiler = profiler
+ patched["read_int"].return_value = SpecialLengths.END_OF_STREAM
+ patched["WorkerPerfProfiler"].return_value = contextlib.nullcontext()
+ patched["WorkerMemoryProfiler"].return_value = contextlib.nullcontext()
+
+ worker_utils.worker_run(main, infile, outfile)
+
+ profiler_calls = (
+ patched["WorkerPerfProfiler"].call_count,
+ patched["WorkerMemoryProfiler"].call_count,
+ )
+
+ outfile.seek(0)
+ return (
+ read_serialized_int(outfile),
+ set(_accumulatorRegistry),
+ profiler_calls,
+ main_call_count,
+ )
+
+ profiler_accumulator_ids = {
+ SpecialAccumulatorIds.SQL_UDF_PROFIER,
+ SpecialAccumulatorIds.SQL_UDF_PROFIER_V2,
+ }
+ test_cases = [
+ (None, 0, set(), (0, 0)),
+ ("unsupported", 0, set(), (0, 0)),
+ ("perf", 2, profiler_accumulator_ids, (1, 0)),
+ ("memory", 2, profiler_accumulator_ids, (0, 1)),
+ ]
+ for profiler, update_count, accumulator_ids, profiler_calls in test_cases:
+ with self.subTest(profiler=profiler):
+ self.assertEqual(
+ run_worker(profiler),
+ (update_count, accumulator_ids, profiler_calls, 1),
+ )
+
+ self.assertEqual(
+ run_worker(None, regular_accumulator_id=1),
+ (1, {1}, (0, 0), 1),
+ )
+
+
@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message)
class BasePythonDataSourceTestsMixin:
spark: SparkSession
diff --git a/python/pyspark/sql/worker/utils.py b/python/pyspark/sql/worker/utils.py
index e22745cc4fc7..2069cdc9e6d0 100644
--- a/python/pyspark/sql/worker/utils.py
+++ b/python/pyspark/sql/worker/utils.py
@@ -110,29 +110,32 @@ def worker_run(main: Callable, infile: IO, outfile: IO) -> None:
conf = RunnerConf(infile)
_accumulatorRegistry.clear()
- accumulator = _deserialize_accumulator(
- SpecialAccumulatorIds.SQL_UDF_PROFIER, {}, ProfileResultsParam
- )
-
- accumulator_v2 = _deserialize_accumulator(
- SpecialAccumulatorIds.SQL_UDF_PROFIER_V2, {}, ProfileResultsParamV2
- )
-
- if main.__module__ == "__main__":
- try:
- worker_module = sys.modules["__main__"].__spec__.name # type: ignore[union-attr]
- except Exception:
- worker_module = "__main__"
- else:
- worker_module = main.__module__
- worker_module = worker_module.split(".")[-1]
-
- if conf.profiler == "perf":
- with WorkerPerfProfiler(accumulator, accumulator_v2, worker_module):
- main(infile, outfile)
- elif conf.profiler == "memory":
- with WorkerMemoryProfiler(accumulator, accumulator_v2, worker_module, main):
- main(infile, outfile)
+ profiler = conf.profiler
+ if profiler in ("perf", "memory"):
+ accumulator = _deserialize_accumulator(
+ SpecialAccumulatorIds.SQL_UDF_PROFIER, {}, ProfileResultsParam
+ )
+ accumulator_v2 = _deserialize_accumulator(
+ SpecialAccumulatorIds.SQL_UDF_PROFIER_V2, {}, ProfileResultsParamV2
+ )
+
+ if main.__module__ == "__main__":
+ try:
+ worker_module = (
+ sys.modules["__main__"].__spec__.name # type: ignore[union-attr]
+ )
+ except Exception:
+ worker_module = "__main__"
+ else:
+ worker_module = main.__module__
+ worker_module = worker_module.split(".")[-1]
+
+ if profiler == "perf":
+ with WorkerPerfProfiler(accumulator, accumulator_v2, worker_module):
+ main(infile, outfile)
+ else:
+ with WorkerMemoryProfiler(accumulator, accumulator_v2, worker_module, main):
+ main(infile, outfile)
else:
main(infile, outfile)
except BaseException as e:
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org