You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by yi...@apache.org on 2026/10/19 16:17:47 UTC
(spark) branch master updated: [SPARK-60280][PYTHON] Replace _is_iter_based with an EvalTypeHandler class property
This is an automated email from the ASF dual-hosted git repository.
Yicong-Huang 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 4deb510051bb [SPARK-60280][PYTHON] Replace _is_iter_based with an EvalTypeHandler class property
4deb510051bb is described below
commit 4deb510051bb7db3a5e8b3ff93804aa7dccaed41
Author: Yicong Huang <17...@users.noreply.github.com>
AuthorDate: Mon Oct 19 16:17:30 2026 +0000
[SPARK-60280][PYTHON] Replace _is_iter_based with an EvalTypeHandler class property
### What changes were proposed in this pull request?
Replace the hardcoded `_is_iter_based(eval_type)` set in `worker.py` with an `is_iter_based` property on the handler metaclass, backed by an `_iter_based` class variable (default `False`) that the handlers of the 9 iterator-based eval types set to `True` (`_ScalarIterElementwiseUDFHandler` covers both ITER_ELEMENTWISE types). `read_single_udf` and `read_single_udtf` now take the resolved `handler_cls` instead of `eval_type` / `partition_wrapper`, and the unused `udf_index` parameter i [...]
### Why are the changes needed?
Keeps per-eval-type facts on the handler instead of a parallel list in the worker. Part of SPARK-59415.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing tests: `pyspark.sql.tests.test_udf_profiler`, `pyspark.tests.test_memory_profiler`, `pyspark.sql.tests.connect.test_parity_udf_profiler`, `pyspark.sql.tests.connect.test_parity_memory_profiler`, `pyspark.sql.tests.test_udtf`, `pyspark.sql.tests.arrow.test_arrow_udtf`, and `pyspark.eval_handlers.tests`.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 5.5
This pull request and its description were written by Isaac.
Closes #59543 from Yicong-Huang/iter-based-handler-property.
Authored-by: Yicong Huang <17...@users.noreply.github.com>
Signed-off-by: Yicong-Huang <17...@users.noreply.github.com>
---
python/pyspark/eval_handlers/_arrow.py | 3 ++
python/pyspark/eval_handlers/_base.py | 9 +++++
python/pyspark/eval_handlers/_elementwise.py | 2 ++
python/pyspark/eval_handlers/_pandas.py | 3 ++
python/pyspark/eval_handlers/_stateful.py | 1 +
python/pyspark/worker.py | 52 ++++++++++------------------
6 files changed, 36 insertions(+), 34 deletions(-)
diff --git a/python/pyspark/eval_handlers/_arrow.py b/python/pyspark/eval_handlers/_arrow.py
index 47eed960672b..2c9acf409235 100644
--- a/python/pyspark/eval_handlers/_arrow.py
+++ b/python/pyspark/eval_handlers/_arrow.py
@@ -130,6 +130,7 @@ class ArrowGroupedMapIterUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
to the declared schema."""
eval_type = PythonEvalType.SQL_GROUPED_MAP_ARROW_ITER_UDF
+ _iter_based = True
def __init__(
self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
@@ -680,6 +681,7 @@ class ArrowMapUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
columns on the wire and wrapped back into a single struct column."""
eval_type = PythonEvalType.SQL_MAP_ARROW_ITER_UDF
+ _iter_based = True
def __init__(
self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
@@ -723,6 +725,7 @@ class ArrowScalarIterUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
result and verify the total row count matches the input."""
eval_type = PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF
+ _iter_based = True
def __init__(
self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
diff --git a/python/pyspark/eval_handlers/_base.py b/python/pyspark/eval_handlers/_base.py
index dde9416f54c7..ce9d2a468a2e 100644
--- a/python/pyspark/eval_handlers/_base.py
+++ b/python/pyspark/eval_handlers/_base.py
@@ -120,6 +120,13 @@ class _EvalTypeHandlerMeta(ABCMeta):
_eval_type_handlers[eval_type] = (conf_name, variants)
return cls
+ _iter_based: bool
+
+ @property
+ def is_iter_based(cls) -> bool:
+ """Whether the handler's UDFs return an iterator."""
+ return cls._iter_based
+
class EvalTypeHandler(Generic[InputBatch, OutputBatch], metaclass=_EvalTypeHandlerMeta):
"""Base class for the UDF execution model.
@@ -138,6 +145,8 @@ class EvalTypeHandler(Generic[InputBatch, OutputBatch], metaclass=_EvalTypeHandl
# the value of it this handler serves; None when the eval type has a single handler.
variant_conf: ClassVar[Optional[str]] = None
variant_value: ClassVar[Any] = None
+ # Whether the handler's UDFs return an iterator.
+ _iter_based: ClassVar[bool] = False
def __init__(
self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf"
diff --git a/python/pyspark/eval_handlers/_elementwise.py b/python/pyspark/eval_handlers/_elementwise.py
index 0b723c6f2f15..cfcae6c81d30 100644
--- a/python/pyspark/eval_handlers/_elementwise.py
+++ b/python/pyspark/eval_handlers/_elementwise.py
@@ -516,6 +516,8 @@ class _ScalarIterElementwiseUDFHandler(_VectorizedElementwiseUDFHandler):
input ones. Subclasses supply the flavor's conversions and result checks.
"""
+ _iter_based = True
+
def __init__(
self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
) -> None:
diff --git a/python/pyspark/eval_handlers/_pandas.py b/python/pyspark/eval_handlers/_pandas.py
index 858836bcdf1a..ba7070657d6b 100644
--- a/python/pyspark/eval_handlers/_pandas.py
+++ b/python/pyspark/eval_handlers/_pandas.py
@@ -149,6 +149,7 @@ class PandasScalarIterUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
Arrow<->pandas conversions in ``pyspark.sql.conversion`` directly."""
eval_type = PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF
+ _iter_based = True
def __init__(
self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf"
@@ -220,6 +221,7 @@ class PandasMapUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
RecordBatch each. The pandas counterpart of ArrowMapUDFHandler."""
eval_type = PythonEvalType.SQL_MAP_PANDAS_ITER_UDF
+ _iter_based = True
def __init__(
self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf"
@@ -369,6 +371,7 @@ class PandasGroupedMapIterUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
pandas counterpart of ArrowGroupedMapIterUDFHandler."""
eval_type = PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF
+ _iter_based = True
def __init__(
self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf"
diff --git a/python/pyspark/eval_handlers/_stateful.py b/python/pyspark/eval_handlers/_stateful.py
index 6ebd7d4625c1..564b25f339fa 100644
--- a/python/pyspark/eval_handlers/_stateful.py
+++ b/python/pyspark/eval_handlers/_stateful.py
@@ -77,6 +77,7 @@ class PandasGroupedMapWithStateUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]
into count/data/state RecordBatches."""
eval_type = PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE
+ _iter_based = True
# The output RecordBatch has three struct fields, accessed by position
# (_0/_1/_2, not by name): a count column indicating how many data and
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 201c52816d46..c76dd0c08f93 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -99,24 +99,7 @@ def chain(f, g):
return lambda *a: g(f(*a))
-def _is_iter_based(eval_type: int) -> bool:
- return eval_type in (
- PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF,
- PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF,
- # Iterator UDFs lifted out of a higher-order function lambda keep the iterator contract:
- # the user function still consumes and produces an iterator of batches; the worker only
- # feeds it the flattened elements and re-nests the results. See ExtractPythonUDFFromLambda.
- PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF,
- PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF,
- PythonEvalType.SQL_MAP_PANDAS_ITER_UDF,
- PythonEvalType.SQL_MAP_ARROW_ITER_UDF,
- PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE,
- PythonEvalType.SQL_GROUPED_MAP_ARROW_ITER_UDF,
- PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF,
- )
-
-
-def wrap_perf_profiler(f, eval_type, result_id):
+def wrap_perf_profiler(f, iter_based, result_id):
from pyspark.sql.profiler import ProfileResultsParam, ProfileResultsParamV2, WorkerPerfProfiler
accumulator = _deserialize_accumulator(
@@ -126,7 +109,7 @@ def wrap_perf_profiler(f, eval_type, result_id):
SpecialAccumulatorIds.SQL_UDF_PROFIER_V2, {}, ProfileResultsParamV2
)
- if _is_iter_based(eval_type):
+ if iter_based:
def profiling_func(*args, **kwargs):
iterator = iter(f(*args, **kwargs))
@@ -148,7 +131,7 @@ def wrap_perf_profiler(f, eval_type, result_id):
return profiling_func
-def wrap_memory_profiler(f, eval_type, result_id):
+def wrap_memory_profiler(f, iter_based, result_id):
import pyspark.memory_profiler_ext
from pyspark.sql.profiler import (
ProfileResultsParam,
@@ -167,7 +150,7 @@ def wrap_memory_profiler(f, eval_type, result_id):
SpecialAccumulatorIds.SQL_UDF_PROFIER_V2, {}, ProfileResultsParamV2
)
- if _is_iter_based(eval_type):
+ if iter_based:
def profiling_func(*args, **kwargs):
g = f(*args, **kwargs)
@@ -191,7 +174,8 @@ def wrap_memory_profiler(f, eval_type, result_id):
return profiling_func
-def read_single_udf(pickleSer, udf_info, eval_type, runner_conf, udf_index):
+def read_single_udf(pickleSer, udf_info, handler_cls, runner_conf):
+ eval_type = handler_cls.eval_type
chained_func = None
for udf in udf_info.udfs:
f, return_type = read_command(pickleSer, udf)
@@ -208,9 +192,13 @@ def read_single_udf(pickleSer, udf_info, eval_type, runner_conf, udf_index):
else:
profiler = runner_conf.udf_profiler
if profiler == "perf":
- profiling_func = wrap_perf_profiler(chained_func, eval_type, udf_info.result_id)
+ profiling_func = wrap_perf_profiler(
+ chained_func, handler_cls.is_iter_based, udf_info.result_id
+ )
elif profiler == "memory":
- profiling_func = wrap_memory_profiler(chained_func, eval_type, udf_info.result_id)
+ profiling_func = wrap_memory_profiler(
+ chained_func, handler_cls.is_iter_based, udf_info.result_id
+ )
else:
profiling_func = chained_func
@@ -300,12 +288,12 @@ def read_single_udf(pickleSer, udf_info, eval_type, runner_conf, udf_index):
# Read a serialized user-defined table function (UDTF) and prepare it for evaluation: wrap
-# the constructor to receive the AnalyzeResult, instantiate the UDTF (wrapped in
+# the constructor to receive the AnalyzeResult, instantiate the UDTF (wrapped in the handler's
# partition_wrapper when its TABLE argument has PARTITION BY expressions), and check the 'eval'
# signature against the call arguments. Returns the same
# (func, args_offsets, kwargs_offsets, return_type) shape as read_single_udf, with the UDTF
# instance in the func slot.
-def read_single_udtf(pickleSer, udtf_info, partition_wrapper):
+def read_single_udtf(pickleSer, udtf_info, handler_cls):
if udtf_info.pickled_analyze_result is not None:
pickled_analyze_result = pickleSer.loads(udtf_info.pickled_analyze_result)
else:
@@ -353,7 +341,7 @@ def read_single_udtf(pickleSer, udtf_info, partition_wrapper):
# Instantiate the UDTF class.
try:
if len(udtf_info.partition_child_indexes) > 0:
- udtf = partition_wrapper(handler, udtf_info.partition_child_indexes)
+ udtf = handler_cls.partition_wrapper(handler, udtf_info.partition_child_indexes)
else:
udtf = handler()
except Exception as e:
@@ -434,16 +422,12 @@ def invoke_udf(message_receiver: SparkMessageReceiver, outfile: BinaryIO):
raise ValueError("Unknown eval type: {}".format(eval_type))
if isinstance(init_info.udf_info, UDTFInfo):
assert issubclass(handler_cls, (UDTFEvalTypeHandler, TableUDTFHandler))
- udfs = [
- read_single_udtf(pickleSer, init_info.udf_info, handler_cls.partition_wrapper)
- ]
+ udfs = [read_single_udtf(pickleSer, init_info.udf_info, handler_cls)]
else:
assert isinstance(init_info.udf_info, list)
udfs = [
- read_single_udf(
- pickleSer, udf_info, eval_type, runner_conf, udf_index=udf_index
- )
- for udf_index, udf_info in enumerate(init_info.udf_info)
+ read_single_udf(pickleSer, udf_info, handler_cls, runner_conf)
+ for udf_info in init_info.udf_info
]
handler = handler_cls(udfs=udfs, runner_conf=runner_conf, eval_conf=eval_conf)
func, serializer = handler.run, handler.serializer
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org