[SPARK-31849][PYTHON][SQL] Make PySpark SQL exceptions more Pythonic

python/pyspark/sql/tests/test_pandas_udf.py

@@ -19,14 +19,12 @@
 
 from pyspark.sql.functions import udf, pandas_udf, PandasUDFType
 from pyspark.sql.types import *
-from pyspark.sql.utils import ParseException
+from pyspark.sql.utils import ParseException, PythonException
 from pyspark.rdd import PythonEvalType
 from pyspark.testing.sqlutils import ReusedSQLTestCase, have_pandas, have_pyarrow, \
     pandas_requirement_message, pyarrow_requirement_message
 from pyspark.testing.utils import QuietTest
 
-from py4j.protocol import Py4JJavaError
-
 
 @unittest.skipIf(
     not have_pandas or not have_pyarrow,
@@ -157,14 +155,14 @@ def foofoo(x, y):
 
         # plain udf (test for SPARK-23754)
         self.assertRaisesRegexp(
-            Py4JJavaError,
+            PythonException,
             exc_message,
             df.withColumn('v', udf(foo)('id')).collect
         )
 
         # pandas scalar udf
         self.assertRaisesRegexp(
-            Py4JJavaError,
+            PythonException,
             exc_message,
             df.withColumn(
                 'v', pandas_udf(foo, 'double', PandasUDFType.SCALAR)('id')
@@ -173,15 +171,15 @@ def foofoo(x, y):
 
         # pandas grouped map
         self.assertRaisesRegexp(
-            Py4JJavaError,
+            PythonException,
             exc_message,
             df.groupBy('id').apply(
                 pandas_udf(foo, df.schema, PandasUDFType.GROUPED_MAP)
             ).collect
         )
 
         self.assertRaisesRegexp(
-            Py4JJavaError,
+            PythonException,
             exc_message,
             df.groupBy('id').apply(
                 pandas_udf(foofoo, df.schema, PandasUDFType.GROUPED_MAP)
@@ -190,7 +188,7 @@ def foofoo(x, y):
 
         # pandas grouped agg
         self.assertRaisesRegexp(
-            Py4JJavaError,
+            PythonException,
             exc_message,
             df.groupBy('id').agg(
                 pandas_udf(foo, 'double', PandasUDFType.GROUPED_AGG)('id')

python/pyspark/sql/utils.py

@@ -18,8 +18,19 @@
 import py4j
 import sys
 
+from pyspark import SparkContext
+
 if sys.version_info.major >= 3:
     unicode = str
+    # Disable exception chaining (PEP 3134) in captured exceptions
+    # in order to hide JVM stacktace.
+    exec("""
+def raise_from(e):
+    raise e from None
+""")
+else:
+    def raise_from(e):
+        raise e
 
 
 class CapturedException(Exception):
@@ -29,7 +40,11 @@ def __init__(self, desc, stackTrace, cause=None):
         self.cause = convert_exception(cause) if cause is not None else None
 
     def __str__(self):
+        sql_conf = SparkContext._jvm.org.apache.spark.sql.internal.SQLConf.get()
+        debug_enabled = sql_conf.pysparkJVMStacktraceEnabled()
         desc = self.desc
+        if debug_enabled:
+            desc = desc + "\nJVM stacktrace:\n%s" % self.stackTrace
         # encode unicode instance for python2 for human readable description
         if sys.version_info.major < 3 and isinstance(desc, unicode):
             return str(desc.encode('utf-8'))
@@ -67,6 +82,12 @@ class QueryExecutionException(CapturedException):
     """
 
 
+class PythonException(CapturedException):
+    """
+    Exceptions thrown from Python workers.
+    """
+
+
 class UnknownException(CapturedException):
     """
     None of the above exceptions.
@@ -75,21 +96,33 @@ class UnknownException(CapturedException):
 
 def convert_exception(e):
     s = e.toString()
-    stackTrace = '\n\t at '.join(map(lambda x: x.toString(), e.getStackTrace()))
     c = e.getCause()
+
+    jvm = SparkContext._jvm
+    jwriter = jvm.java.io.StringWriter()
+    e.printStackTrace(jvm.java.io.PrintWriter(jwriter))
+    stacktrace = jwriter.toString()
     if s.startswith('org.apache.spark.sql.AnalysisException: '):
-        return AnalysisException(s.split(': ', 1)[1], stackTrace, c)
+        return AnalysisException(s.split(': ', 1)[1], stacktrace, c)
     if s.startswith('org.apache.spark.sql.catalyst.analysis'):
-        return AnalysisException(s.split(': ', 1)[1], stackTrace, c)
+        return AnalysisException(s.split(': ', 1)[1], stacktrace, c)
     if s.startswith('org.apache.spark.sql.catalyst.parser.ParseException: '):
-        return ParseException(s.split(': ', 1)[1], stackTrace, c)
+        return ParseException(s.split(': ', 1)[1], stacktrace, c)
     if s.startswith('org.apache.spark.sql.streaming.StreamingQueryException: '):
-        return StreamingQueryException(s.split(': ', 1)[1], stackTrace, c)
+        return StreamingQueryException(s.split(': ', 1)[1], stacktrace, c)
     if s.startswith('org.apache.spark.sql.execution.QueryExecutionException: '):
-        return QueryExecutionException(s.split(': ', 1)[1], stackTrace, c)
+        return QueryExecutionException(s.split(': ', 1)[1], stacktrace, c)
     if s.startswith('java.lang.IllegalArgumentException: '):
-        return IllegalArgumentException(s.split(': ', 1)[1], stackTrace, c)
-    return UnknownException(s, stackTrace, c)
+        return IllegalArgumentException(s.split(': ', 1)[1], stacktrace, c)
+    if c is not None and (
+            c.toString().startswith('org.apache.spark.api.python.PythonException: ')
+            # To make sure this only catches Python UDFs.
+            and any(map(lambda v: "org.apache.spark.sql.execution.python" in v.toString(),
+                        c.getStackTrace()))):
+        msg = ("\n  An exception was thrown from Python worker in the executor. "
+               "The below is the Python worker stacktrace.\n%s" % c.getMessage())
+        return PythonException(msg, stacktrace)
+    return UnknownException(s, stacktrace, c)
 
 
 def capture_sql_exception(f):
@@ -99,7 +132,9 @@ def deco(*a, **kw):
         except py4j.protocol.Py4JJavaError as e:
             converted = convert_exception(e.java_exception)
             if not isinstance(converted, UnknownException):
-                raise converted
+                # Hide where the exception came from that shows a non-Pythonic
+                # JVM exception message.
+                raise_from(converted)
             else:
                 raise
     return deco

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -1784,6 +1784,15 @@ object SQLConf {
       .version("3.0.0")
       .fallbackConf(ARROW_EXECUTION_ENABLED)
 
+  val PYSPARK_JVM_STACKTRACE_ENABLED =
+    buildConf("spark.sql.pyspark.jvmStacktrace.enabled")
+      .doc("When true, it shows the JVM stacktrace in the user-facing PySpark exception " +
+        "together with Python stacktrace. By default, it is disabled and hides JVM stacktrace " +
+        "and shows a Python-friendly exception only.")
+      .version("3.0.0")
+      .booleanConf
+      .createWithDefault(false)
+
   val ARROW_SPARKR_EXECUTION_ENABLED =
     buildConf("spark.sql.execution.arrow.sparkr.enabled")
       .doc("When true, make use of Apache Arrow for columnar data transfers in SparkR. " +
@@ -3063,6 +3072,8 @@ class SQLConf extends Serializable with Logging {
 
   def arrowPySparkEnabled: Boolean = getConf(ARROW_PYSPARK_EXECUTION_ENABLED)
 
+  def pysparkJVMStacktraceEnabled: Boolean = getConf(PYSPARK_JVM_STACKTRACE_ENABLED)
+
   def arrowSparkREnabled: Boolean = getConf(ARROW_SPARKR_EXECUTION_ENABLED)
 
   def arrowPySparkFallbackEnabled: Boolean = getConf(ARROW_PYSPARK_FALLBACK_ENABLED)