This is an automated email from the ASF dual-hosted git repository.
HyukjinKwon pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new cc76e900cabb [SPARK-56519][PYTHON] Isolate communication part from
python udf worker
cc76e900cabb is described below
commit cc76e900cabb2682edebc5f7c7d0f02e2feaad1b
Author: Tian Gao <[email protected]>
AuthorDate: Fri May 8 08:45:47 2026 +0900
[SPARK-56519][PYTHON] Isolate communication part from python udf worker
### What changes were proposed in this pull request?
We read everything we need for creating Python UDF worker from JVM at the
very beginning, so we don't need to read anything while we process. This
isolates the communication part from the processing logic.
A message interface is designed to hold all the information from JVM. For
now it does not change the protocol at all. It is designed to support multiple
protocols. The goal is to have a layer that contains all the information needed
in a language agnostic form (`int`, `str`, `bytes`, `None`, `list`, `dict`). In
the future, we can support a whole different protocol while keep supporting the
existing socket protocol - the new protocol just needs to fill in the blanks.
We still have some random workers that rely on `setup_spark_files` or
`setup_broadcasts`. So in this PR we made all these worker utils support both
the old and new input. If the argument is an IO, it uses the original code
path. If the function is fed the concrete data, it just uses it to initialize
stuff.
### Why are the changes needed?
The current protocol is super fragile because we mix the communication and
processing. No one knows where exactly the socket is going to be read. It's
also heavily based on the actual eval type which is horrible. We want to
isolate the communication part out and do read just once - read everything we
need then process. It will make the protocol much more robust and easier to
debug.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
`test_udf.py` and `test_udtf.py` passed locally. Still need to wait for CI
result.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #55380 from gaogaotiantian/isolate-init-data.
Authored-by: Tian Gao <[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>
(cherry picked from commit 0fc754675ed96b025b0a7f4953600b627925bdd3)
Signed-off-by: Hyukjin Kwon <[email protected]>
---
python/pyspark/taskcontext.py | 6 +-
python/pyspark/worker.py | 126 +++++++-----------
python/pyspark/worker_message.py | 272 +++++++++++++++++++++++++++++++++++++++
python/pyspark/worker_util.py | 145 +++++++++++++--------
4 files changed, 414 insertions(+), 135 deletions(-)
diff --git a/python/pyspark/taskcontext.py b/python/pyspark/taskcontext.py
index 8347e77cb94e..7ac645323243 100644
--- a/python/pyspark/taskcontext.py
+++ b/python/pyspark/taskcontext.py
@@ -14,7 +14,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
-from typing import ClassVar, Type, TypeVar, Dict, List, Optional, Union, cast
+from typing import Any, ClassVar, Type, TypeVar, Dict, List, Optional, Union,
cast
from pyspark.util import local_connect_and_auth
from pyspark.serializers import read_int, write_int, write_with_length,
UTF8Deserializer
@@ -130,7 +130,7 @@ class TaskContext:
_cpus: Optional[int] = None
_resources: Optional[Dict[str, "ResourceInformation"]] = None
- def __new__(cls: Type["TaskContext"], **kwargs: Dict) -> "TaskContext":
+ def __new__(cls: Type["TaskContext"], **kwargs: Any) -> "TaskContext":
"""
Even if users construct :class:`TaskContext` instead of using get,
give them the singleton.
"""
@@ -142,7 +142,7 @@ class TaskContext:
def __init__(
self,
- **kwargs: Dict,
+ **kwargs: Any,
) -> None:
# Set attributes only if they are passed in and not None
# The kwargs are auto-mapped to the private attributes of TaskContext
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index a47e940f031f..b9f791ada1e6 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -57,8 +57,6 @@ from pyspark.taskcontext import BarrierTaskContext,
TaskContext
from pyspark.util import PythonEvalType
from pyspark.serializers import (
write_int,
- read_long,
- read_bool,
write_long,
read_int,
SpecialLengths,
@@ -109,6 +107,7 @@ from pyspark.util import (
)
from pyspark import _NoValue, shuffle
from pyspark.errors import PySparkRuntimeError, PySparkTypeError,
PySparkValueError
+from pyspark.worker_message import WorkerInitInfo
from pyspark.worker_util import (
check_python_version,
get_sock_file_to_executor,
@@ -118,7 +117,6 @@ from pyspark.worker_util import (
setup_broadcasts,
setup_memory_limits,
setup_spark_files,
- utf8_deserializer,
Conf,
)
from pyspark.logger.worker_io import capture_outputs
@@ -1017,29 +1015,15 @@ def wrap_memory_profiler(f, eval_type, result_id):
return profiling_func
-def read_single_udf(pickleSer, infile, eval_type, runner_conf, udf_index):
- num_arg = read_int(infile)
-
- args_offsets = []
- kwargs_offsets = {}
- for _ in range(num_arg):
- offset = read_int(infile)
- if read_bool(infile):
- name = utf8_deserializer.loads(infile)
- kwargs_offsets[name] = offset
- else:
- args_offsets.append(offset)
-
+def read_single_udf(pickleSer, udf_info, eval_type, runner_conf, udf_index):
chained_func = None
- for i in range(read_int(infile)):
- f, return_type = read_command(pickleSer, infile)
+ for udf in udf_info.udfs:
+ f, return_type = read_command(pickleSer, udf)
if chained_func is None:
chained_func = f
else:
chained_func = chain(chained_func, f)
- result_id = read_long(infile)
-
# If chained_func is from pyspark.sql.worker, it is to read/write data
source.
# In this case, we check the data_source_profiler config.
if getattr(chained_func, "__module__",
"").startswith("pyspark.sql.worker."):
@@ -1047,9 +1031,9 @@ def read_single_udf(pickleSer, infile, eval_type,
runner_conf, udf_index):
else:
profiler = runner_conf.udf_profiler
if profiler == "perf":
- profiling_func = wrap_perf_profiler(chained_func, eval_type, result_id)
+ profiling_func = wrap_perf_profiler(chained_func, eval_type,
udf_info.result_id)
elif profiler == "memory":
- profiling_func = wrap_memory_profiler(chained_func, eval_type,
result_id)
+ profiling_func = wrap_memory_profiler(chained_func, eval_type,
udf_info.result_id)
else:
profiling_func = chained_func
@@ -1063,6 +1047,8 @@ def read_single_udf(pickleSer, infile, eval_type,
runner_conf, udf_index):
# when they are processed in a for loop, raise them as RuntimeError's
instead
func = fail_on_stopiteration(profiling_func)
+ args_offsets, kwargs_offsets = udf_info.args, udf_info.kwargs
+
# the last returnType will be the return type of UDF
if eval_type == PythonEvalType.SQL_SCALAR_PANDAS_UDF:
return wrap_scalar_pandas_udf(func, args_offsets, kwargs_offsets,
return_type, runner_conf)
@@ -1141,7 +1127,7 @@ def read_single_udf(pickleSer, infile, eval_type,
runner_conf, udf_index):
# It expects the UDTF to be in a specific format and performs various checks to
# ensure the UDTF is valid. This function also prepares a mapper function for
applying
# the UDTF logic to input rows.
-def read_udtf(pickleSer, infile, eval_type, runner_conf, eval_conf):
+def read_udtf(pickleSer, udtf_info, eval_type, runner_conf, eval_conf):
if eval_type == PythonEvalType.SQL_ARROW_TABLE_UDF:
if runner_conf.use_legacy_pandas_udtf_conversion:
# NOTE: if timezone is set here, that implies
respectSessionTimeZone is True
@@ -1162,50 +1148,35 @@ def read_udtf(pickleSer, infile, eval_type,
runner_conf, eval_conf):
# Each row is a group so do not batch but send one by one.
ser = BatchedSerializer(CPickleSerializer(), 1)
- # See 'PythonUDTFRunner.PythonUDFWriterThread.writeCommand'
- num_arg = read_int(infile)
- args_offsets = []
- kwargs_offsets = {}
- for _ in range(num_arg):
- offset = read_int(infile)
- if read_bool(infile):
- name = utf8_deserializer.loads(infile)
- kwargs_offsets[name] = offset
- else:
- args_offsets.append(offset)
- num_partition_child_indexes = read_int(infile)
- partition_child_indexes = [read_int(infile) for i in
range(num_partition_child_indexes)]
- has_pickled_analyze_result = read_bool(infile)
- if has_pickled_analyze_result:
- pickled_analyze_result = pickleSer._read_with_length(infile)
+ if udtf_info.pickled_analyze_result is not None:
+ pickled_analyze_result =
pickleSer.loads(udtf_info.pickled_analyze_result)
else:
pickled_analyze_result = None
# Initially we assume that the UDTF __init__ method accepts the pickled
AnalyzeResult,
# although we may set this to false later if we find otherwise.
- handler = read_command(pickleSer, infile)
+ handler = read_command(pickleSer, udtf_info.handler)
if not isinstance(handler, type):
raise PySparkRuntimeError(
f"Invalid UDTF handler type. Expected a class (type 'type'), but "
f"got an instance of {type(handler).__name__}."
)
- return_type = _parse_datatype_json_string(utf8_deserializer.loads(infile))
+ return_type = _parse_datatype_json_string(udtf_info.return_type)
if not isinstance(return_type, StructType):
raise PySparkRuntimeError(
f"The return type of a UDTF must be a struct type, but got
{type(return_type)}."
)
- udtf_name = utf8_deserializer.loads(infile)
# Update the handler that creates a new UDTF instance to first try calling
the UDTF constructor
# with one argument containing the previous AnalyzeResult. If that fails,
then try a constructor
# with no arguments. In this way each UDTF class instance can decide if it
wants to inspect the
# AnalyzeResult.
udtf_init_args = inspect.getfullargspec(handler)
- if has_pickled_analyze_result:
+ if pickled_analyze_result is not None:
if len(udtf_init_args.args) > 2:
raise PySparkRuntimeError(
errorClass="UDTF_CONSTRUCTOR_INVALID_IMPLEMENTS_ANALYZE_METHOD",
- messageParameters={"name": udtf_name},
+ messageParameters={"name": udtf_info.name},
)
elif len(udtf_init_args.args) == 2:
prev_handler = handler
@@ -1218,7 +1189,7 @@ def read_udtf(pickleSer, infile, eval_type, runner_conf,
eval_conf):
elif len(udtf_init_args.args) > 1:
raise PySparkRuntimeError(
errorClass="UDTF_CONSTRUCTOR_INVALID_NO_ANALYZE_METHOD",
- messageParameters={"name": udtf_name},
+ messageParameters={"name": udtf_info.name},
)
class UDTFWithPartitions:
@@ -1256,7 +1227,7 @@ def read_udtf(pickleSer, infile, eval_type, runner_conf,
eval_conf):
self._create_udtf: Callable = create_udtf
self._udtf = create_udtf()
self._prev_arguments: list = list()
- self._partition_child_indexes: list = partition_child_indexes
+ self._partition_child_indexes: list =
udtf_info.partition_child_indexes
self._eval_raised_skip_rest_of_input_table: bool = False
def eval(self, *args, **kwargs) -> Iterator:
@@ -1650,13 +1621,13 @@ def read_udtf(pickleSer, infile, eval_type,
runner_conf, eval_conf):
# Instantiate the UDTF class.
try:
- if len(partition_child_indexes) > 0:
+ if len(udtf_info.partition_child_indexes) > 0:
# Determine if this is an Arrow UDTF
is_arrow_udtf = eval_type == PythonEvalType.SQL_ARROW_UDTF
if is_arrow_udtf:
- udtf = ArrowUDTFWithPartition(handler, partition_child_indexes)
+ udtf = ArrowUDTFWithPartition(handler,
udtf_info.partition_child_indexes)
else:
- udtf = UDTFWithPartitions(handler, partition_child_indexes)
+ udtf = UDTFWithPartitions(handler,
udtf_info.partition_child_indexes)
else:
udtf = handler()
except Exception as e:
@@ -1676,11 +1647,11 @@ def read_udtf(pickleSer, infile, eval_type,
runner_conf, eval_conf):
# Check that the arguments provided to the UDTF call match the expected
parameters defined
# in the 'eval' method signature.
try:
- inspect.signature(udtf.eval).bind(*args_offsets, **kwargs_offsets)
+ inspect.signature(udtf.eval).bind(*udtf_info.args, **udtf_info.kwargs)
except TypeError as e:
raise PySparkRuntimeError(
errorClass="UDTF_EVAL_METHOD_ARGUMENTS_DO_NOT_MATCH_SIGNATURE",
- messageParameters={"name": udtf_name, "reason": str(e)},
+ messageParameters={"name": udtf_info.name, "reason": str(e)},
) from None
def build_null_checker(return_type: StructType) ->
Optional[Callable[[Any], None]]:
@@ -1884,7 +1855,7 @@ def read_udtf(pickleSer, infile, eval_type, runner_conf,
eval_conf):
return evaluate
eval_func_kwargs_support, args_kwargs_offsets = wrap_kwargs_support(
- getattr(udtf, "eval"), args_offsets, kwargs_offsets
+ getattr(udtf, "eval"), udtf_info.args, udtf_info.kwargs
)
eval = wrap_arrow_udtf(eval_func_kwargs_support, return_type)
@@ -2046,7 +2017,7 @@ def read_udtf(pickleSer, infile, eval_type, runner_conf,
eval_conf):
return evaluate
eval_func_kwargs_support, args_kwargs_offsets = wrap_kwargs_support(
- getattr(udtf, "eval"), args_offsets, kwargs_offsets
+ getattr(udtf, "eval"), udtf_info.args, udtf_info.kwargs
)
eval = wrap_arrow_udtf(eval_func_kwargs_support, return_type)
@@ -2172,7 +2143,7 @@ def read_udtf(pickleSer, infile, eval_type, runner_conf,
eval_conf):
return evaluate
eval_func_kwargs_support, args_kwargs_offsets = wrap_kwargs_support(
- getattr(udtf, "eval"), args_offsets, kwargs_offsets
+ getattr(udtf, "eval"), udtf_info.args, udtf_info.kwargs
)
eval = wrap_pyarrow_udtf(eval_func_kwargs_support, return_type)
@@ -2274,7 +2245,7 @@ def read_udtf(pickleSer, infile, eval_type, runner_conf,
eval_conf):
return evaluate
eval_func_kwargs_support, args_kwargs_offsets = wrap_kwargs_support(
- getattr(udtf, "eval"), args_offsets, kwargs_offsets
+ getattr(udtf, "eval"), udtf_info.args, udtf_info.kwargs
)
eval = wrap_udtf(eval_func_kwargs_support, return_type)
@@ -2302,7 +2273,7 @@ def read_udtf(pickleSer, infile, eval_type, runner_conf,
eval_conf):
return mapper, None, ser, ser
-def read_udfs(pickleSer, infile, eval_type, runner_conf, eval_conf):
+def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
if eval_type in (
PythonEvalType.SQL_ARROW_BATCHED_UDF,
PythonEvalType.SQL_SCALAR_PANDAS_UDF,
@@ -2436,13 +2407,13 @@ def read_udfs(pickleSer, infile, eval_type,
runner_conf, eval_conf):
batch_size = int(os.environ.get("PYTHON_UDF_BATCH_SIZE", "100"))
ser = BatchedSerializer(CPickleSerializer(), batch_size)
- # Read all UDFs
- num_udfs = read_int(infile)
udfs = [
- read_single_udf(pickleSer, infile, eval_type, runner_conf, udf_index=i)
- for i in range(num_udfs)
+ read_single_udf(pickleSer, udf_info, eval_type, runner_conf,
udf_index=udf_index)
+ for udf_index, udf_info in enumerate(udf_info_list)
]
+ num_udfs = len(udfs)
+
def extract_key_value_indexes(grouped_arg_offsets):
"""
Helper function to extract the key and value indexes from arg_offsets
for the grouped and
@@ -3555,52 +3526,49 @@ def read_udfs(pickleSer, infile, eval_type,
runner_conf, eval_conf):
def main(infile, outfile):
try:
boot_time = time.time()
- split_index = read_int(infile)
- if split_index == -1: # for unit tests
- sys.exit(-1)
+ init_info = WorkerInitInfo.from_stream(infile)
start_faulthandler_periodic_traceback()
- check_python_version(infile)
+ check_python_version(init_info.python_version)
memory_limit_mb = int(os.environ.get("PYSPARK_EXECUTOR_MEMORY_MB",
"-1"))
setup_memory_limits(memory_limit_mb)
- task_context_json = json.loads(utf8_deserializer.loads(infile))
- if task_context_json["isBarrier"]:
- taskContext = BarrierTaskContext.from_json(task_context_json)
- else:
- taskContext = TaskContext.from_json(task_context_json)
- TaskContext._setTaskContext(taskContext)
+ TaskContext._setTaskContext(init_info.task_context.to_task_context())
shuffle.MemoryBytesSpilled = 0
shuffle.DiskBytesSpilled = 0
- setup_spark_files(infile)
- setup_broadcasts(infile)
+ setup_spark_files(init_info.spark_files_dir, init_info.python_includes)
+ setup_broadcasts(
+ init_info.broadcast.variables,
+ init_info.broadcast.conn_info,
+ init_info.broadcast.auth_secret,
+ )
_accumulatorRegistry.clear()
- eval_type = read_int(infile)
- runner_conf = RunnerConf(infile)
- eval_conf = EvalConf(infile)
+ eval_type = init_info.eval_type
+ runner_conf = RunnerConf(init_info.runner_conf)
+ eval_conf = EvalConf(init_info.eval_conf)
if eval_type == PythonEvalType.NON_UDF:
- func, profiler, deserializer, serializer = read_command(pickleSer,
infile)
+ func, profiler, deserializer, serializer = read_command(pickleSer,
init_info.udf_info)
elif eval_type in (
PythonEvalType.SQL_TABLE_UDF,
PythonEvalType.SQL_ARROW_TABLE_UDF,
PythonEvalType.SQL_ARROW_UDTF,
):
func, profiler, deserializer, serializer = read_udtf(
- pickleSer, infile, eval_type, runner_conf, eval_conf
+ pickleSer, init_info.udf_info, eval_type, runner_conf,
eval_conf
)
else:
func, profiler, deserializer, serializer = read_udfs(
- pickleSer, infile, eval_type, runner_conf, eval_conf
+ pickleSer, init_info.udf_info, eval_type, runner_conf,
eval_conf
)
init_time = time.time()
def process():
iterator = deserializer.load_stream(infile)
- out_iter = func(split_index, iterator)
+ out_iter = func(init_info.split_index, iterator)
try:
serializer.dump_stream(out_iter, outfile)
finally:
diff --git a/python/pyspark/worker_message.py b/python/pyspark/worker_message.py
new file mode 100644
index 000000000000..b1519cb08442
--- /dev/null
+++ b/python/pyspark/worker_message.py
@@ -0,0 +1,272 @@
+#
+# 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 dataclasses
+import json
+import sys
+from typing import Optional, Union, IO
+
+from pyspark.errors import PySparkValueError
+from pyspark.serializers import read_bool, read_int, read_long, SpecialLengths
+from pyspark.taskcontext import BarrierTaskContext, ResourceInformation,
TaskContext
+from pyspark.util import PythonEvalType
+from pyspark.worker_util import utf8_deserializer
+
+
[email protected]
+class TaskContextInfo:
+ @dataclasses.dataclass
+ class ResourceInfo:
+ name: str
+ addresses: list[str]
+
+ is_barrier: bool
+ conn_info: Optional[Union[str, int]]
+ secret: Optional[str]
+ stage_id: int
+ partition_id: int
+ attempt_number: int
+ task_attempt_id: int
+ cpus: int
+ resources: dict[str, ResourceInfo]
+ local_properties: dict[str, str]
+
+ @classmethod
+ def from_stream(cls, stream: IO) -> "TaskContextInfo":
+ task_context_json = json.loads(utf8_deserializer.loads(stream))
+ return cls(
+ is_barrier=task_context_json["isBarrier"],
+ conn_info=task_context_json["connInfo"],
+ secret=task_context_json["secret"],
+ stage_id=task_context_json["stageId"],
+ partition_id=task_context_json["partitionId"],
+ attempt_number=task_context_json["attemptNumber"],
+ task_attempt_id=task_context_json["taskAttemptId"],
+ cpus=task_context_json["cpus"],
+ resources={
+ k: cls.ResourceInfo(name=v["name"], addresses=v["addresses"])
+ for k, v in task_context_json["resources"].items()
+ },
+ local_properties=task_context_json["localProperties"],
+ )
+
+ def to_task_context(self) -> TaskContext:
+ if self.is_barrier:
+ return BarrierTaskContext(
+ conn_info=self.conn_info,
+ secret=self.secret,
+ stageId=self.stage_id,
+ partitionId=self.partition_id,
+ attemptNumber=self.attempt_number,
+ taskAttemptId=self.task_attempt_id,
+ cpus=self.cpus,
+ resources={
+ k: ResourceInformation(v.name, v.addresses) for k, v in
self.resources.items()
+ },
+ localProperties=self.local_properties,
+ )
+ else:
+ return TaskContext(
+ stageId=self.stage_id,
+ partitionId=self.partition_id,
+ attemptNumber=self.attempt_number,
+ taskAttemptId=self.task_attempt_id,
+ cpus=self.cpus,
+ resources={
+ k: ResourceInformation(v.name, v.addresses) for k, v in
self.resources.items()
+ },
+ localProperties=self.local_properties,
+ )
+
+
[email protected]
+class BroadcastInfo:
+ conn_info: Optional[Union[str, int]]
+ auth_secret: Optional[str]
+ variables: list[tuple[int, Optional[str]]]
+
+ @classmethod
+ def from_stream(cls, stream: IO) -> "BroadcastInfo":
+ needs_broadcast_decryption_server = read_bool(stream)
+ num_broadcast_variables = read_int(stream)
+ conn_info = None
+ auth_secret = None
+ if needs_broadcast_decryption_server:
+ conn_info = read_int(stream)
+ if conn_info == -1:
+ conn_info = utf8_deserializer.loads(stream)
+ else:
+ auth_secret = utf8_deserializer.loads(stream)
+
+ variables = []
+ for _ in range(num_broadcast_variables):
+ bid = read_long(stream)
+ path = None
+ if bid >= 0 and not needs_broadcast_decryption_server:
+ path = utf8_deserializer.loads(stream)
+ variables.append((bid, path))
+
+ return cls(conn_info=conn_info, auth_secret=auth_secret,
variables=variables)
+
+
[email protected]
+class UDFInfo:
+ udfs: list[bytes]
+ args: list[int]
+ kwargs: dict[str, int]
+ result_id: int
+
+ @classmethod
+ def from_stream(cls, stream: IO) -> "UDFInfo":
+ num_args = read_int(stream)
+ udfs = []
+ args = []
+ kwargs = {}
+
+ for _ in range(num_args):
+ offset = read_int(stream)
+ if read_bool(stream):
+ name = utf8_deserializer.loads(stream)
+ kwargs[name] = offset
+ else:
+ args.append(offset)
+
+ for i in range(read_int(stream)):
+ length = read_int(stream)
+ if length == SpecialLengths.END_OF_DATA_SECTION:
+ raise EOFError
+ elif length == SpecialLengths.NULL:
+ raise PySparkValueError("Unexpected NULL value for UDF")
+ else:
+ data = stream.read(length)
+ if len(data) < length:
+ raise EOFError
+ udfs.append(data)
+
+ result_id = read_long(stream)
+
+ return cls(udfs=udfs, args=args, kwargs=kwargs, result_id=result_id)
+
+
[email protected]
+class UDTFInfo:
+ args: list[int]
+ kwargs: dict[str, int]
+ partition_child_indexes: list[int]
+ pickled_analyze_result: Optional[bytes]
+ handler: bytes
+ return_type: str
+ name: str
+
+ @classmethod
+ def from_stream(cls, stream: IO) -> "UDTFInfo":
+ # See 'PythonUDTFRunner.PythonUDFWriterThread.writeCommand'
+ args = []
+ kwargs = {}
+ for _ in range(read_int(stream)):
+ offset = read_int(stream)
+ if read_bool(stream):
+ name = utf8_deserializer.loads(stream)
+ kwargs[name] = offset
+ else:
+ args.append(offset)
+ partition_child_indexes = [read_int(stream) for _ in
range(read_int(stream))]
+ if read_bool(stream):
+ pickled_analyze_result = stream.read(read_int(stream))
+ else:
+ pickled_analyze_result = None
+ handler = stream.read(read_int(stream))
+ return_type = utf8_deserializer.loads(stream)
+ name = utf8_deserializer.loads(stream)
+
+ return cls(
+ args=args,
+ kwargs=kwargs,
+ partition_child_indexes=partition_child_indexes,
+ pickled_analyze_result=pickled_analyze_result,
+ handler=handler,
+ return_type=return_type,
+ name=name,
+ )
+
+
[email protected]
+class WorkerInitInfo:
+ split_index: int
+ python_version: str
+ spark_files_dir: str
+ task_context: TaskContextInfo
+ python_includes: list[str]
+ broadcast: BroadcastInfo
+ eval_type: int
+ runner_conf: dict[str, str]
+ eval_conf: dict[str, str]
+ udf_info: Union[bytes, UDTFInfo, list[UDFInfo]]
+
+ @classmethod
+ def from_stream(cls, stream: IO) -> "WorkerInitInfo":
+ split_index = read_int(stream)
+ if split_index == -1:
+ sys.exit(-1)
+ python_version = utf8_deserializer.loads(stream)
+ task_context = TaskContextInfo.from_stream(stream)
+
+ spark_files_dir = utf8_deserializer.loads(stream)
+ python_includes = []
+ for _ in range(read_int(stream)):
+ python_includes.append(utf8_deserializer.loads(stream))
+
+ broadcast = BroadcastInfo.from_stream(stream)
+ eval_type = read_int(stream)
+ runner_conf = {}
+ for _ in range(read_int(stream)):
+ k = utf8_deserializer.loads(stream)
+ v = utf8_deserializer.loads(stream)
+ runner_conf[k] = v
+ eval_conf = {}
+ for _ in range(read_int(stream)):
+ k = utf8_deserializer.loads(stream)
+ v = utf8_deserializer.loads(stream)
+ eval_conf[k] = v
+
+ udf_info: Union[bytes, UDTFInfo, list[UDFInfo]]
+
+ if eval_type == PythonEvalType.NON_UDF:
+ udf_info = stream.read(read_int(stream))
+ elif eval_type in (
+ PythonEvalType.SQL_TABLE_UDF,
+ PythonEvalType.SQL_ARROW_TABLE_UDF,
+ PythonEvalType.SQL_ARROW_UDTF,
+ ):
+ udf_info = UDTFInfo.from_stream(stream)
+ else:
+ udf_info = []
+ for _ in range(read_int(stream)):
+ udf_info.append(UDFInfo.from_stream(stream))
+
+ return cls(
+ split_index=split_index,
+ python_version=python_version,
+ spark_files_dir=spark_files_dir,
+ task_context=task_context,
+ python_includes=python_includes,
+ broadcast=broadcast,
+ eval_type=eval_type,
+ runner_conf=runner_conf,
+ eval_conf=eval_conf,
+ udf_info=udf_info,
+ )
diff --git a/python/pyspark/worker_util.py b/python/pyspark/worker_util.py
index 8c52bde6ea6c..08edf0c5decb 100644
--- a/python/pyspark/worker_util.py
+++ b/python/pyspark/worker_util.py
@@ -24,7 +24,7 @@ import importlib
from inspect import currentframe, getframeinfo
import os
import sys
-from typing import Any, Generator, IO, Optional
+from typing import Any, Generator, IO, Optional, Union, overload
import warnings
if "SPARK_TESTING" in os.environ:
@@ -44,7 +44,6 @@ from pyspark.util import is_remote_only
from pyspark.errors import PySparkRuntimeError
from pyspark.util import local_connect_and_auth
from pyspark.serializers import (
- read_bool,
read_int,
read_long,
write_int,
@@ -66,21 +65,27 @@ def add_path(path: str) -> bool:
return False
-def read_command(serializer: FramedSerializer, file: IO) -> Any:
+def read_command(serializer: FramedSerializer, file: Union[IO, bytes]) -> Any:
if not is_remote_only():
from pyspark.core.broadcast import Broadcast
- command = serializer._read_with_length(file)
+ if isinstance(file, bytes):
+ command = serializer.loads(file)
+ else:
+ command = serializer._read_with_length(file)
if not is_remote_only() and isinstance(command, Broadcast):
command = serializer.loads(command.value)
return command
-def check_python_version(infile: IO) -> None:
+def check_python_version(infile_or_version: Union[IO, str]) -> None:
"""
Check the Python version between the running process and the one used to
serialize the command.
"""
- version = utf8_deserializer.loads(infile)
+ if isinstance(infile_or_version, str):
+ version = infile_or_version
+ else:
+ version = utf8_deserializer.loads(infile_or_version)
worker_version = "%d.%d" % sys.version_info[:2]
if version != worker_version:
raise PySparkRuntimeError(
@@ -130,12 +135,20 @@ def setup_memory_limits(memory_limit_mb: int) -> None:
)
-def setup_spark_files(infile: IO) -> None:
+@overload
+def setup_spark_files(infile_or_spark_files_dir: IO) -> None: ...
+@overload
+def setup_spark_files(infile_or_spark_files_dir: str, python_includes:
list[str]) -> None: ...
+def setup_spark_files(
+ infile_or_spark_files_dir: Union[IO, str], python_includes:
Optional[list[str]] = None
+) -> None:
"""
Set up Spark files, archives, and pyfiles.
"""
- # fetch name of workdir
- spark_files_dir = utf8_deserializer.loads(infile)
+ if isinstance(infile_or_spark_files_dir, str):
+ spark_files_dir = infile_or_spark_files_dir
+ else:
+ spark_files_dir = utf8_deserializer.loads(infile_or_spark_files_dir)
if not is_remote_only():
from pyspark.core.files import SparkFiles
@@ -145,51 +158,74 @@ def setup_spark_files(infile: IO) -> None:
# fetch names of includes (*.zip and *.egg files) and construct PYTHONPATH
path_changed = add_path(spark_files_dir) # *.py files that were added
will be copied here
- num_python_includes = read_int(infile)
- for _ in range(num_python_includes):
- filename = utf8_deserializer.loads(infile)
+ if not isinstance(infile_or_spark_files_dir, str):
+ python_includes = [
+ utf8_deserializer.loads(infile_or_spark_files_dir)
+ for _ in range(read_int(infile_or_spark_files_dir))
+ ]
+ assert python_includes is not None
+
+ for filename in python_includes:
path_changed = add_path(os.path.join(spark_files_dir, filename)) or
path_changed
if path_changed:
importlib.invalidate_caches()
-def setup_broadcasts(infile: IO) -> None:
+@overload
+def setup_broadcasts(infile_or_variables: IO) -> None: ...
+@overload
+def setup_broadcasts(
+ infile_or_variables: list[tuple[int, Union[str, None]]], conn_info: str,
auth_secret: None
+) -> None: ...
+@overload
+def setup_broadcasts(
+ infile_or_variables: list[tuple[int, Union[str, None]]], conn_info: int,
auth_secret: str
+) -> None: ...
+@overload
+def setup_broadcasts(
+ infile_or_variables: list[tuple[int, Union[str, None]]], conn_info: None,
auth_secret: None
+) -> None: ...
+def setup_broadcasts(
+ infile_or_variables: Union[IO, list[tuple[int, Union[str, None]]]],
+ conn_info: Optional[Union[str, int]] = None,
+ auth_secret: Optional[str] = None,
+) -> None:
"""
Set up broadcasted variables.
"""
if not is_remote_only():
from pyspark.core.broadcast import Broadcast, _broadcastRegistry
- # fetch names and values of broadcast variables
- needs_broadcast_decryption_server = read_bool(infile)
- num_broadcast_variables = read_int(infile)
+ if isinstance(infile_or_variables, list):
+ variables = infile_or_variables
+ else:
+ from pyspark.worker_message import BroadcastInfo
+
+ broadcast_info = BroadcastInfo.from_stream(infile_or_variables)
+ conn_info = broadcast_info.conn_info
+ auth_secret = broadcast_info.auth_secret
+ variables = broadcast_info.variables
+
+ needs_broadcast_decryption_server = conn_info is not None or auth_secret
is not None
+
if needs_broadcast_decryption_server:
- # read the decrypted data from a server in the jvm
- conn_info = read_int(infile)
- auth_secret = None
- if conn_info == -1:
- conn_info = utf8_deserializer.loads(infile)
- else:
- auth_secret = utf8_deserializer.loads(infile)
broadcast_sock_file, _ = local_connect_and_auth(conn_info, auth_secret)
+ else:
+ broadcast_sock_file = None
- for _ in range(num_broadcast_variables):
- bid = read_long(infile)
+ for bid, path in variables:
if bid >= 0:
- if needs_broadcast_decryption_server:
+ if path is None:
read_bid = read_long(broadcast_sock_file)
assert read_bid == bid
_broadcastRegistry[bid] =
Broadcast(sock_file=broadcast_sock_file)
else:
- path = utf8_deserializer.loads(infile)
_broadcastRegistry[bid] = Broadcast(path=path)
-
else:
- bid = -bid - 1
- _broadcastRegistry.pop(bid)
+ _broadcastRegistry.pop(-bid - 1)
- if needs_broadcast_decryption_server:
+ if broadcast_sock_file is not None:
broadcast_sock_file.write(b"1")
broadcast_sock_file.close()
@@ -223,29 +259,32 @@ def send_accumulator_updates(outfile: IO) -> None:
class Conf:
- def __init__(self, infile: Optional[IO] = None) -> None:
+ def __init__(self, infile_or_dict: Optional[Union[dict[str, str], IO]] =
None) -> None:
self._conf: dict[str, Any] = {}
- if infile is not None:
- self.load(infile)
-
- def load(self, infile: IO) -> None:
- num_conf = read_int(infile)
- # We do a sanity check here to reduce the possibility to stuck
indefinitely
- # due to an invalid messsage. If the numer of configurations is
obviously
- # wrong, we just raise an error directly.
- # We hand-pick the configurations to send to the worker so the number
should
- # be very small (less than 100).
- if num_conf < 0 or num_conf > 10000:
- raise PySparkRuntimeError(
- errorClass="PROTOCOL_ERROR",
- messageParameters={
- "failure": f"Invalid number of configurations: {num_conf}",
- },
- )
- for _ in range(num_conf):
- k = utf8_deserializer.loads(infile)
- v = utf8_deserializer.loads(infile)
- self._conf[k] = v
+ if infile_or_dict is not None:
+ self.load(infile_or_dict)
+
+ def load(self, infile_or_dict: Union[dict[str, str], IO]) -> None:
+ if isinstance(infile_or_dict, dict):
+ self._conf = infile_or_dict
+ else:
+ num_conf = read_int(infile_or_dict)
+ # We do a sanity check here to reduce the possibility to stuck
indefinitely
+ # due to an invalid messsage. If the numer of configurations is
obviously
+ # wrong, we just raise an error directly.
+ # We hand-pick the configurations to send to the worker so the
number should
+ # be very small (less than 100).
+ if num_conf < 0 or num_conf > 10000:
+ raise PySparkRuntimeError(
+ errorClass="PROTOCOL_ERROR",
+ messageParameters={
+ "failure": f"Invalid number of configurations:
{num_conf}",
+ },
+ )
+ for _ in range(num_conf):
+ k = utf8_deserializer.loads(infile_or_dict)
+ v = utf8_deserializer.loads(infile_or_dict)
+ self._conf[k] = v
def get(self, key: str, default: Any = "", *, lower_str: bool = True) ->
Any:
val = self._conf.get(key, default)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
