[SPARK-58575][PYTHON] Inline redundant error_class parameter in worker.py verify helpers

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())

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