You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by yi...@apache.org on 2026/08/04 19:09:50 UTC
(spark) branch branch-4.x updated: [SPARK-58529][PYTHON] Unify RESULT_ROWS_MISMATCH message and consolidate row count verification in worker.py
This is an automated email from the ASF dual-hosted git repository.
Yicong-Huang pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new c2c21fef09b4 [SPARK-58529][PYTHON] Unify RESULT_ROWS_MISMATCH message and consolidate row count verification in worker.py
c2c21fef09b4 is described below
commit c2c21fef09b46b0e9d6e5743c131bcedc7fabf03
Author: Yicong Huang <17...@users.noreply.github.com>
AuthorDate: Tue Aug 4 19:09:26 2026 +0000
[SPARK-58529][PYTHON] Unify RESULT_ROWS_MISMATCH message and consolidate row count verification in worker.py
### What changes were proposed in this pull request?
`worker.py` raises `RESULT_ROWS_MISMATCH` from four places on a UDF row-count mismatch. SPARK-58161 hardcoded a `message=` override at each, carrying a pandas-specific tail (`Result vector from pandas_udf was not the required length...`) that is wrong on the arrow UDF path (`verify_scalar_result` also serves arrow UDFs). This PR removes all four overrides so the generic `RESULT_ROWS_MISMATCH` template applies, and consolidates the two row-count helpers: the iterator generator (renamed [...]
### Why are the changes needed?
The pandas-specific tail is misleading for arrow UDFs, and the mismatch check was duplicated across four sites. One helper plus the generic template removes the duplication and makes the message consistent and correct.
### Does this PR introduce _any_ user-facing change?
Yes. The `RESULT_ROWS_MISMATCH` message drops the trailing `Result vector from pandas_udf was not the required length: expected N, got M.`; the leading `The number of output rows (M) must match the number of input rows (N).` and the error class / `messageParameters` are unchanged.
### How was this patch tested?
Existing assertions in `test_pandas_udf_scalar.py` already match the template text and still pass. Helpers were also exercised in isolation across list, iterator, callable, and scalar paths.
Note: `pyspark-connect-old-client` fails because branch-4.0's `test_vectorized_udf_invalid_length` asserts on the removed substring. This is the intended message change, not a regression; a branch-4.0 backport will relax that assertion after merge.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #57728 from Yicong-Huang/merge-verify-rowcount.
Authored-by: Yicong Huang <17...@users.noreply.github.com>
Signed-off-by: Yicong-Huang <17...@users.noreply.github.com>
(cherry picked from commit 44ea3d172bba43f1fec5f3791c0d2470e5d13a86)
Signed-off-by: Yicong-Huang <17...@users.noreply.github.com>
---
python/pyspark/worker.py | 63 ++++++++++--------------------------------------
1 file changed, 13 insertions(+), 50 deletions(-)
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 0db7e7151038..5dc22650455c 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,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):
@@ -2043,10 +2017,9 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
)
# 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 +3043,7 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
"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 +3133,9 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
)
# 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
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org