[SPARK-XXXXX][PYTHON] Standardize PySpark row count mismatch error class to use RESULT_ROWS_MISMATCH

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:
 

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 <udf_type> was not the required length: expected <expected>, got <actual>."
-    ]
-  },
+
   "SESSION_ALREADY_EXIST": {
     "message": [
       "Cannot start a remote Spark session because there is a regular Spark session already running."

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()
 

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(