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 <[email protected]>
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 <[email protected]>
Signed-off-by: Yicong-Huang <[email protected]>
(cherry picked from commit 44ea3d172bba43f1fec5f3791c0d2470e5d13a86)
Signed-off-by: Yicong-Huang <[email protected]>
---
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: [email protected]
For additional commands, e-mail: [email protected]
