[SPARK-50960][PYTHON][CONNECT] Add InvalidPlanInput to Spark Connec…

…t Python client ### What changes were proposed in this pull request? This PR proposes to add `InvalidPlanInput` to Spark Connect Python client. ### Why are the changes needed? To keep the consistency with Spark Connect Scala client. We should capture the `SparkConnectGrpcException` more generally across all clients. ### Does this PR introduce _any_ user-facing change? No API changes, but the user-facing error message would keep consistency with Spark Connect Scala client. ### How was this patch tested? Updated the existing UT. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #49616 from itholic/invalid_input. Authored-by: Haejoon Lee <[email protected]> Signed-off-by: Hyukjin Kwon <[email protected]>
apache · Feb 13, 2025 · f749742 · f749742
1 parent dc56829
commit f749742
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/python/pyspark/errors/exceptions/connect.py b/python/pyspark/errors/exceptions/connect.py
@@ -295,6 +295,12 @@ class SparkNoSuchElementException(SparkConnectGrpcException, BaseNoSuchElementEx
     """
 
 
+class InvalidPlanInput(SparkConnectGrpcException):
+    """
+    Error thrown when a connect plan is not valid.
+    """
+
+
 # Update EXCEPTION_CLASS_MAPPING here when adding a new exception
 EXCEPTION_CLASS_MAPPING = {
     "org.apache.spark.sql.catalyst.parser.ParseException": ParseException,
@@ -312,6 +318,7 @@ class SparkNoSuchElementException(SparkConnectGrpcException, BaseNoSuchElementEx
     "org.apache.spark.api.python.PythonException": PythonException,
     "org.apache.spark.SparkNoSuchElementException": SparkNoSuchElementException,
     "org.apache.spark.SparkException": SparkException,
+    "org.apache.spark.sql.connect.common.InvalidPlanInput": InvalidPlanInput,
 }
 
 

diff --git a/python/pyspark/sql/tests/connect/test_parity_udtf.py b/python/pyspark/sql/tests/connect/test_parity_udtf.py
@@ -26,7 +26,11 @@
 
     sql.udtf.UserDefinedTableFunction = UserDefinedTableFunction
     from pyspark.sql.connect.functions import lit, udtf
-    from pyspark.errors.exceptions.connect import SparkConnectGrpcException, PythonException
+    from pyspark.errors.exceptions.connect import (
+        SparkConnectGrpcException,
+        PythonException,
+        InvalidPlanInput,
+    )
 
 
 class UDTFParityTests(BaseUDTFTestsMixin, ReusedConnectTestCase):
@@ -54,7 +58,7 @@ def eval(self, a: int):
                 yield a + 1,
 
         with self.assertRaisesRegex(
-            SparkConnectGrpcException, "Invalid Python user-defined table function return type."
+            InvalidPlanInput, "Invalid Python user-defined table function return type."
         ):
             TestUDTF(lit(1)).collect()