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 c6c238106b48 [SPARK-60529][PYTHON] Wrap pandas scalar iterator UDFs in
fail_on_stopiteration
c6c238106b48 is described below
commit c6c238106b489969f66d9935d98226db56862d3f
Author: Yicong Huang <[email protected]>
AuthorDate: Tue Oct 27 08:05:49 2026 +0000
[SPARK-60529][PYTHON] Wrap pandas scalar iterator UDFs in
fail_on_stopiteration
### What changes were proposed in this pull request?
Enable the default `fail_on_stopiteration` wrap for
`SQL_SCALAR_PANDAS_ITER_UDF` (`PandasScalarIterUDFHandler`) and
`SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF`
(`PandasScalarIterElementwiseUDFHandler`), matching their Arrow counterparts,
and correct the `_fail_on_stopiteration` comment: since `run` is a generator, a
leaked StopIteration already fails the task via PEP 479; the wrap only turns it
into a clear `STOP_ITERATION_OCCURRED`. It also makes the iterator UDF
profilers call the UDF [...]
### Why are the changes needed?
Consistent error across the scalar iterator handlers; the old comment
overstated the wrap's effect.
### Does this PR introduce _any_ user-facing change?
Yes, small. A StopIteration raised by the UDF call itself (before it
returns an iterator) now fails with `PySparkRuntimeError
[STOP_ITERATION_OCCURRED]` instead of `RuntimeError: generator raised
StopIteration`. StopIteration raised while iterating the returned iterator is
unchanged.
### How was this patch tested?
New unit tests in `test_pandas_eval_type_handlers.py` /
`test_elementwise_eval_type_handlers.py`; new
`test_iter_profiler_keeps_stop_iteration_error`; updated
`test_fail_on_stopiteration_per_handler`.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Opus 5.5
This pull request and its description were written by Isaac.
Closes #59774 from Yicong-Huang/SPARK-60529.
Authored-by: Yicong Huang <[email protected]>
Signed-off-by: Yicong-Huang <[email protected]>
---
python/pyspark/eval_handlers/_base.py | 46 +++++++++++++---------
python/pyspark/eval_handlers/_elementwise.py | 1 -
python/pyspark/eval_handlers/_pandas.py | 1 -
.../tests/test_base_eval_type_handlers.py | 20 +++++++---
.../tests/test_elementwise_eval_type_handlers.py | 9 +++++
.../tests/test_pandas_eval_type_handlers.py | 11 ++++++
6 files changed, 63 insertions(+), 25 deletions(-)
diff --git a/python/pyspark/eval_handlers/_base.py
b/python/pyspark/eval_handlers/_base.py
index d9bc7354b2b2..327f5189c4cd 100644
--- a/python/pyspark/eval_handlers/_base.py
+++ b/python/pyspark/eval_handlers/_base.py
@@ -150,8 +150,8 @@ class EvalTypeHandler(Generic[InputBatch, OutputBatch],
metaclass=_EvalTypeHandl
variant_value: ClassVar[Any] = None
# Whether the handler's UDFs return an iterator.
_iter_based: ClassVar[bool] = False
- # Whether __init__ wraps the UDF in fail_on_stopiteration, so a
StopIteration raised in
- # user code fails the task instead of silently ending the loop that calls
the UDF.
+ # Whether __init__ wraps the UDF in fail_on_stopiteration, so a
StopIteration raised by
+ # the UDF call fails with STOP_ITERATION_OCCURRED instead of PEP 479's
generic RuntimeError.
_fail_on_stopiteration: ClassVar[bool] = True
def __init__(
@@ -207,14 +207,20 @@ class EvalTypeHandler(Generic[InputBatch, OutputBatch],
metaclass=_EvalTypeHandl
if self._iter_based:
def profiling_func(*args: Any, **kwargs: Any) -> Any:
+ # Call f eagerly, as the unprofiled path does, so
fail_on_stopiteration sees a
+ # StopIteration raised by the call; profile each next() lazily.
iterator = iter(f(*args, **kwargs))
- while True:
- try:
- with WorkerPerfProfiler(accumulator, accumulator_v2,
result_id):
- item = next(iterator)
- yield item
- except StopIteration:
- break
+
+ def profiled() -> Iterator[Any]:
+ while True:
+ try:
+ with WorkerPerfProfiler(accumulator,
accumulator_v2, result_id):
+ item = next(iterator)
+ yield item
+ except StopIteration:
+ break
+
+ return profiled()
else:
@@ -248,18 +254,22 @@ class EvalTypeHandler(Generic[InputBatch, OutputBatch],
metaclass=_EvalTypeHandl
if self._iter_based:
def profiling_func(*args: Any, **kwargs: Any) -> Any:
+ # Call f eagerly, as in the perf profiler.
g = f(*args, **kwargs)
iterator = iter(g)
- while True:
- try:
- with WorkerMemoryProfiler(
- accumulator, accumulator_v2, result_id, g.gi_code
- ):
- item = next(iterator)
- yield item
- except StopIteration:
- break
+ def profiled() -> Iterator[Any]:
+ while True:
+ try:
+ with WorkerMemoryProfiler(
+ accumulator, accumulator_v2, result_id,
g.gi_code
+ ):
+ item = next(iterator)
+ yield item
+ except StopIteration:
+ break
+
+ return profiled()
else:
diff --git a/python/pyspark/eval_handlers/_elementwise.py
b/python/pyspark/eval_handlers/_elementwise.py
index 66d5994039b4..7646b8254026 100644
--- a/python/pyspark/eval_handlers/_elementwise.py
+++ b/python/pyspark/eval_handlers/_elementwise.py
@@ -652,7 +652,6 @@ class PandasScalarIterElementwiseUDFHandler(
an iterator of pandas Series (DataFrames for struct types)."""
eval_type = PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF
- _fail_on_stopiteration = False
def _verify_return_type(self, results: Any) -> Iterator[Any]:
import pandas as pd
diff --git a/python/pyspark/eval_handlers/_pandas.py
b/python/pyspark/eval_handlers/_pandas.py
index 9a503ee425ab..a9215e95bb21 100644
--- a/python/pyspark/eval_handlers/_pandas.py
+++ b/python/pyspark/eval_handlers/_pandas.py
@@ -150,7 +150,6 @@ class
PandasScalarIterUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF
_iter_based = True
- _fail_on_stopiteration = False
def __init__(self, udfs: list[UDF], runner_conf: "RunnerConf", eval_conf:
"EvalConf") -> None:
import pandas as pd
diff --git a/python/pyspark/eval_handlers/tests/test_base_eval_type_handlers.py
b/python/pyspark/eval_handlers/tests/test_base_eval_type_handlers.py
index ba611ba866ab..baff8a51e6d7 100644
--- a/python/pyspark/eval_handlers/tests/test_base_eval_type_handlers.py
+++ b/python/pyspark/eval_handlers/tests/test_base_eval_type_handlers.py
@@ -310,11 +310,7 @@ class HandlerUDFsTests(PySparkBaseTestCase):
udf.raw_func(1)
def test_fail_on_stopiteration_per_handler(self):
- not_wrapped = {
- PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF,
- PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF,
- PythonEvalType.SQL_ARROW_BATCHED_UDF,
- }
+ not_wrapped = {PythonEvalType.SQL_ARROW_BATCHED_UDF}
for _, variants in list(_eval_type_handlers.values()):
for handler_cls in variants.values():
with self.subTest(handler=handler_cls.__name__):
@@ -360,6 +356,20 @@ class HandlerUDFsTests(PySparkBaseTestCase):
self.assertEqual(udf_memory._with_profiler(f, 0), "memory")
self.assertIs(no_profiler._with_profiler(f, 0), f)
+ def test_iter_profiler_keeps_stop_iteration_error(self):
+ # A StopIteration raised by the call of an iterator UDF fails with
+ # STOP_ITERATION_OCCURRED whether or not a profiler is set.
+ class _IterHandler(self._Handler):
+ _iter_based = True
+
+ for profiler in (None, "perf", "memory"):
+ with self.subTest(profiler=profiler):
+ conf = {"spark.sql.pyspark.udf.profiler": profiler} if
profiler else {}
+ handler = self._new(_IterHandler, RunnerConf(conf))
+ with self.assertRaises(PySparkRuntimeError) as cm:
+ handler._funcs[0](1)
+ self.assertEqual(cm.exception.getCondition(),
"STOP_ITERATION_OCCURRED")
+
class UnpickleAndChainTests(PySparkBaseTestCase):
def test_chains_in_order(self):
diff --git
a/python/pyspark/eval_handlers/tests/test_elementwise_eval_type_handlers.py
b/python/pyspark/eval_handlers/tests/test_elementwise_eval_type_handlers.py
index 2b018efc46c7..69a2301e7a64 100644
--- a/python/pyspark/eval_handlers/tests/test_elementwise_eval_type_handlers.py
+++ b/python/pyspark/eval_handlers/tests/test_elementwise_eval_type_handlers.py
@@ -492,6 +492,15 @@ class ScalarIterElementwiseUDFHandlerTestsMixin:
)
)
+ def test_stop_iteration_in_udf_call_raises_runtime_error(self):
+ # A non-generator UDF that eagerly reads an empty input raises
StopIteration from the
+ # call itself.
+ def eager(batches):
+ return iter([next(batches)])
+
+ with self.assertRaisesRegex(PySparkRuntimeError,
"STOP_ITERATION_OCCURRED"):
+ list(self._iter_handler(eager).run(0, iter([])))
+
def test_rejects_non_iterator_result(self):
with self.assertRaisesRegex(PySparkTypeError, "UDF_RETURN_TYPE"):
list(self._iter_handler(lambda batches: [1]).run(0,
iter([self._batch()])))
diff --git
a/python/pyspark/eval_handlers/tests/test_pandas_eval_type_handlers.py
b/python/pyspark/eval_handlers/tests/test_pandas_eval_type_handlers.py
index 2f419d87b0f1..72c101ff901f 100644
--- a/python/pyspark/eval_handlers/tests/test_pandas_eval_type_handlers.py
+++ b/python/pyspark/eval_handlers/tests/test_pandas_eval_type_handlers.py
@@ -378,6 +378,17 @@ class
PandasScalarIterUDFHandlerTests(PandasEvalTypeHandlerTestCase):
list(handler.run(0, iter([self._batch(a=[1, 2]), self._batch(a=[3,
4])])))
self.assertEqual(cm.exception.getCondition(),
"INPUT_NOT_FULLY_CONSUMED")
+ def test_stop_iteration_in_udf_call_raises_runtime_error(self):
+ # A non-generator UDF that eagerly reads an empty input raises
StopIteration from the
+ # call itself.
+ def eager(series_iter):
+ return iter([next(series_iter)])
+
+ handler = self._handler(eager)
+ with self.assertRaises(PySparkRuntimeError) as cm:
+ list(handler.run(0, iter([])))
+ self.assertEqual(cm.exception.getCondition(),
"STOP_ITERATION_OCCURRED")
+
def test_struct_return_takes_dataframe(self):
# struct_in_pandas="dict" + df_for_struct=True: a struct return is a
DataFrame.
struct_type = StructType([StructField("x", LongType()),
StructField("y", LongType())])
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
