Skip to content

[SPARK-58529][PYTHON] Unify RESULT_ROWS_MISMATCH message and consolidate row count verification in worker.py #57728

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Yicong-Huang wants to merge 2 commits into from
+13 −50
Closed

[SPARK-58529][PYTHON] Unify RESULT_ROWS_MISMATCH message and consolidate row count verification in worker.py #57728

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
bc377de
refactor: consolidate UDF row count verification into verify_result_r…
Yicong-Huang Aug 3, 2026
bd0ff65
refactor: narrow verify_iter_result_row_count expected_rows to Callable
Yicong-Huang Aug 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 13 additions & 50 deletions python/pyspark/worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}.",
)


Expand All @@ -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


Expand Down Expand Up @@ -364,37 +353,22 @@ 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,
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
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_rows())


def wrap_udf(f, args_offsets, kwargs_offsets, return_type):
Expand Down Expand Up @@ -2043,10 +2017,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
Expand Down Expand Up @@ -3070,16 +3043,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
Expand Down Expand Up @@ -3169,10 +3133,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
Expand Down