From bc377dee392368f7bd820d2853976404935466d3 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Mon, 3 Aug 2026 21:50:00 +0000 Subject: [PATCH 1/2] refactor: consolidate UDF row count verification into verify_result_row_count --- python/pyspark/worker.py | 53 +++++----------------------------------- 1 file changed, 6 insertions(+), 47 deletions(-) diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py index 0db7e71510382..18ae6a2f7be1c 100644 --- a/python/pyspark/worker.py +++ b/python/pyspark/worker.py @@ -301,8 +301,6 @@ def verify_result_row_count(result_length: int, expected: int) -> None: "output_length": str(result_length), "input_length": str(expected), }, - message=f"The number of output rows ({result_length}) must match the number of input rows ({expected}). " - f"Result vector from pandas_udf was not the required length: expected {expected}, got {result_length}.", ) @@ -327,16 +325,7 @@ def verify_scalar_result(result: Any, num_rows: int) -> Any: "actual": type(result).__name__, }, ) - if result_length != num_rows: - raise PySparkRuntimeError( - errorClass="RESULT_ROWS_MISMATCH", - messageParameters={ - "output_length": str(result_length), - "input_length": str(num_rows), - }, - message=f"The number of output rows ({result_length}) must match the number of input rows ({num_rows}). " - f"Result vector from pandas_udf was not the required length: expected {num_rows}, got {result_length}.", - ) + verify_result_row_count(result_length, num_rows) return result @@ -364,10 +353,9 @@ def verify_output_row_limit( yield element -def verify_output_row_count( +def verify_iter_result_row_count( iterator: Iterator, expected_rows: Union[int, Callable[[], int]], - error_class: str, ) -> Iterator: """Yield elements and verify final row count matches expected exactly.""" actual_rows = 0 @@ -376,25 +364,7 @@ def verify_output_row_count( yield element expected = expected_rows() if callable(expected_rows) else expected_rows - if actual_rows != expected: - if error_class == "RESULT_ROWS_MISMATCH": - raise PySparkRuntimeError( - errorClass=error_class, - messageParameters={ - "output_length": str(actual_rows), - "input_length": str(expected), - }, - message=f"The number of output rows ({actual_rows}) must match the number of input rows ({expected}). " - f"Result vector from pandas_udf was not the required length: expected {expected}, got {actual_rows}.", - ) - else: - raise PySparkRuntimeError( - errorClass=error_class, - messageParameters={ - "output_length": str(actual_rows), - "input_length": str(expected), - }, - ) + verify_result_row_count(actual_rows, expected) def wrap_udf(f, args_offsets, kwargs_offsets, return_type): @@ -2043,10 +2013,9 @@ def process_results(): ) # Apply row count match check (final) - matched = verify_output_row_count( + matched = verify_iter_result_row_count( limited, lambda: num_input_rows, - error_class="RESULT_ROWS_MISMATCH", ) # Yield batches @@ -3070,16 +3039,7 @@ def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.Record "actual": type(result).__name__, }, ) - if len(result) != num_rows: - raise PySparkRuntimeError( - errorClass="RESULT_ROWS_MISMATCH", - messageParameters={ - "output_length": str(len(result)), - "input_length": str(num_rows), - }, - message=f"The number of output rows ({len(result)}) must match the number of input rows ({num_rows}). " - f"Result vector from pandas_udf was not the required length: expected {num_rows}, got {len(result)}.", - ) + verify_result_row_count(len(result), num_rows) # struct_in_pandas="dict": UDF must return DataFrame for struct types if isinstance(udf_return_type, StructType) and not isinstance( result, pd.DataFrame @@ -3169,10 +3129,9 @@ def process_results(): ) # Apply row count match check (final) - matched = verify_output_row_count( + matched = verify_iter_result_row_count( limited, lambda: num_input_rows, - error_class="RESULT_ROWS_MISMATCH", ) # Yield batches From bd0ff65c8503963945b3b29bcf650fa13932c1b1 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Mon, 3 Aug 2026 22:21:57 +0000 Subject: [PATCH 2/2] refactor: narrow verify_iter_result_row_count expected_rows to Callable --- python/pyspark/worker.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py index 18ae6a2f7be1c..5dc22650455ce 100644 --- a/python/pyspark/worker.py +++ b/python/pyspark/worker.py @@ -355,16 +355,20 @@ def verify_output_row_limit( def verify_iter_result_row_count( iterator: Iterator, - expected_rows: Union[int, Callable[[], int]], + expected_rows: Callable[[], int], ) -> Iterator: - """Yield elements and verify final row count matches expected exactly.""" + """Yield elements and verify final row count matches expected exactly. + + ``expected_rows`` is a callable because the expected count is only known once + the iterator is fully consumed (input rows are counted lazily as a side effect + of pulling batches), so it must be read after this generator is exhausted. + """ actual_rows = 0 for element in iterator: actual_rows += len(element) yield element - expected = expected_rows() if callable(expected_rows) else expected_rows - verify_result_row_count(actual_rows, expected) + verify_result_row_count(actual_rows, expected_rows()) def wrap_udf(f, args_offsets, kwargs_offsets, return_type):