diff --git a/python/docs/source/development/debugging.rst b/python/docs/source/development/debugging.rst index 4b7e2e288301c..e69055b410a21 100644 --- a/python/docs/source/development/debugging.rst +++ b/python/docs/source/development/debugging.rst @@ -575,7 +575,7 @@ Solution: 1 6 dtype: int64 -**RuntimeError: Result vector from pandas_udf was not the required length** +**PySparkRuntimeError: [RESULT_ROWS_MISMATCH] The number of output rows must match the number of input rows** Exception: @@ -588,7 +588,7 @@ Exception: 22/04/12 13:46:39 ERROR Executor: Exception in task 2.0 in stage 16.0 (TID 88) org.apache.spark.api.python.PythonException: Traceback (most recent call last): ... - RuntimeError: Result vector from pandas_udf was not the required length: expected 1, got 0 + pyspark.errors.exceptions.base.PySparkRuntimeError: [RESULT_ROWS_MISMATCH] The number of output rows (0) must match the number of input rows (1). Solution: diff --git a/python/pyspark/errors/error-conditions.json b/python/pyspark/errors/error-conditions.json index 38417cbf01889..0482ab8db8f8d 100644 --- a/python/pyspark/errors/error-conditions.json +++ b/python/pyspark/errors/error-conditions.json @@ -851,11 +851,7 @@ "An Observation can be used with a DataFrame only once." ] }, - "SCHEMA_MISMATCH_FOR_PANDAS_UDF": { - "message": [ - "Result vector from was not the required length: expected , got ." - ] - }, + "SESSION_ALREADY_EXIST": { "message": [ "Cannot start a remote Spark session because there is a regular Spark session already running." diff --git a/python/pyspark/sql/tests/pandas/test_pandas_udf_scalar.py b/python/pyspark/sql/tests/pandas/test_pandas_udf_scalar.py index 72d9fa566deef..4fa4571c69989 100644 --- a/python/pyspark/sql/tests/pandas/test_pandas_udf_scalar.py +++ b/python/pyspark/sql/tests/pandas/test_pandas_udf_scalar.py @@ -678,7 +678,7 @@ def check_vectorized_udf_invalid_length(self): df = self.spark.range(10) raise_exception = pandas_udf(lambda _: pd.Series(1), LongType()) with self.assertRaisesRegex( - Exception, "Result vector from pandas_udf was not the required length" + Exception, "The number of output rows.*must match the number of input rows" ): df.select(raise_exception(col("id"))).collect() diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py index 9e629138d5370..61dc30697df05 100644 --- a/python/pyspark/worker.py +++ b/python/pyspark/worker.py @@ -304,6 +304,8 @@ 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}.", ) @@ -329,14 +331,14 @@ def verify_scalar_result(result: Any, num_rows: int) -> Any: }, ) if result_length != num_rows: - # TODO: change error class to RESULT_ROWS_MISMATCH raise PySparkRuntimeError( - errorClass="SCHEMA_MISMATCH_FOR_PANDAS_UDF", + errorClass="RESULT_ROWS_MISMATCH", messageParameters={ - "udf_type": "arrow_udf", - "expected": str(num_rows), - "actual": str(result_length), + "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}.", ) return result @@ -378,13 +380,24 @@ def verify_output_row_count( expected = expected_rows() if callable(expected_rows) else expected_rows if actual_rows != expected: - raise PySparkRuntimeError( - errorClass=error_class, - messageParameters={ - "output_length": str(actual_rows), - "input_length": str(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), + }, + ) def wrap_udf(f, args_offsets, kwargs_offsets, return_type): @@ -3239,12 +3252,13 @@ def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.Record ) if len(result) != num_rows: raise PySparkRuntimeError( - errorClass="SCHEMA_MISMATCH_FOR_PANDAS_UDF", + errorClass="RESULT_ROWS_MISMATCH", messageParameters={ - "udf_type": "pandas_udf", - "expected": str(num_rows), - "actual": str(len(result)), + "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)}.", ) # struct_in_pandas="dict": UDF must return DataFrame for struct types if isinstance(udf_return_type, StructType) and not isinstance(