diff --git a/python/pyspark/sql/tests/pandas/test_pandas_udf_scalar.py b/python/pyspark/sql/tests/pandas/test_pandas_udf_scalar.py index 4fa4571c69989..326dd9d537261 100644 --- a/python/pyspark/sql/tests/pandas/test_pandas_udf_scalar.py +++ b/python/pyspark/sql/tests/pandas/test_pandas_udf_scalar.py @@ -704,6 +704,18 @@ def iter_udf_not_reading_all_input(it): with self.assertRaisesRegex(Exception, "The input iterator must be fully consumed"): df1.select(iter_udf_not_reading_all_input(col("id"))).collect() + @pandas_udf(LongType(), PandasUDFType.SCALAR_ITER) + def iter_udf_too_many_output_rows(it): + for batch in it: + yield pd.Series([1] * (len(batch) + 1)) + + with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": 3}): + df1 = self.spark.range(10).repartition(1) + with self.assertRaisesRegex( + Exception, "The number of output rows must not exceed the number of input rows" + ): + df1.select(iter_udf_too_many_output_rows(col("id"))).collect() + def test_vectorized_udf_chained(self): df = self.spark.range(10) scalar_f = pandas_udf(lambda x: x + 1, LongType()) diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py index 5dc22650455ce..2bd2ea1a3685a 100644 --- a/python/pyspark/worker.py +++ b/python/pyspark/worker.py @@ -329,27 +329,26 @@ def verify_scalar_result(result: Any, num_rows: int) -> Any: return result -def verify_iterator_exhausted(iterator: Iterator, error_class: str) -> None: +def verify_iterator_exhausted(iterator: Iterator) -> None: """Verify that an iterator has been fully consumed.""" try: next(iterator) except StopIteration: pass else: - raise PySparkRuntimeError(errorClass=error_class, messageParameters={}) + raise PySparkRuntimeError(errorClass="INPUT_NOT_FULLY_CONSUMED", messageParameters={}) def verify_output_row_limit( iterator: Iterator, max_rows: Union[int, Callable[[], int]], - error_class: str, ) -> Iterator: """Yield elements while verifying total rows do not exceed a limit (fail-fast).""" total_rows = 0 for element in iterator: total_rows += len(element) if total_rows > (max_rows() if callable(max_rows) else max_rows): - raise PySparkRuntimeError(errorClass=error_class, messageParameters={}) + raise PySparkRuntimeError(errorClass="OUTPUT_EXCEEDS_INPUT_ROWS", messageParameters={}) yield element @@ -2013,7 +2012,6 @@ def process_results(): limited = verify_output_row_limit( process_results(), lambda: num_input_rows, - error_class="OUTPUT_EXCEEDS_INPUT_ROWS", ) # Apply row count match check (final) @@ -2026,10 +2024,7 @@ def process_results(): yield from matched # Verify iterator consumed - verify_iterator_exhausted( - args_iter, - error_class="INPUT_NOT_FULLY_CONSUMED", - ) + verify_iterator_exhausted(args_iter) # profiling is not supported for UDF return func, None, ser, ser @@ -3129,7 +3124,6 @@ def process_results(): limited = verify_output_row_limit( process_results(), lambda: num_input_rows, - error_class="OUTPUT_EXCEEDS_INPUT_ROWS", ) # Apply row count match check (final) @@ -3142,10 +3136,7 @@ def process_results(): yield from matched # Verify iterator consumed - verify_iterator_exhausted( - args_iter, - error_class="INPUT_NOT_FULLY_CONSUMED", - ) + verify_iterator_exhausted(args_iter) # profiling is not supported for UDF return func, None, ser, ser