This is an automated email from the ASF dual-hosted git repository.
gaogaotiantian 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 a9209c5c0ca9 [SPARK-57271][PYTHON] Add capability to print locals in
traceback for Python UDF
a9209c5c0ca9 is described below
commit a9209c5c0ca9285fc359986427ed3210a64db11a
Author: Tian Gao <[email protected]>
AuthorDate: Fri Jun 12 13:11:21 2026 -0700
[SPARK-57271][PYTHON] Add capability to print locals in traceback for
Python UDF
### What changes were proposed in this pull request?
A new configuration
`spark.sql.execution.pyspark.udf.tracebackWithLocals.enabled` is introduced for
UDF tracebacks to include local variable values.
```
pyspark.errors.exceptions.captured.PythonException: An exception was thrown
from the Python worker:
Traceback (most recent call last):
File "/Users/tian.gao/programs/spark/python/example.py", line 10, in f
return g(y)
^^^^
x = 1
y = 0
File "/Users/tian.gao/programs/spark/python/example.py", line 6, in g
return 1 / y
~~^~~
y = 0
```
### Why are the changes needed?
This is very helpful for UDF debugging - it includes all the local
variables in each frame for users to understand the crash site. This is off by
default so the original behavior is kept.
### Does this PR introduce _any_ user-facing change?
Yes, opt-in configuration to enable extra information in the UDF traceback.
### How was this patch tested?
Some local manual tests were done. Two test cases were added. Pending CI.
### Was this patch authored or co-authored using generative AI tooling?
Tests were generated by Claude Code (Opus 4.8).
Closes #56336 from gaogaotiantian/exception-with-locals.
Authored-by: Tian Gao <[email protected]>
Signed-off-by: Tian Gao <[email protected]>
(cherry picked from commit bef8152d6fb1df0f8eb34808411dcd61125f18cc)
Signed-off-by: Tian Gao <[email protected]>
---
.../org/apache/spark/api/python/PythonRunner.scala | 4 ++++
python/pyspark/sql/tests/test_udf.py | 14 ++++++++++++++
python/pyspark/tests/test_util.py | 20 +++++++++++++++++++-
python/pyspark/util.py | 15 +++++++++++----
.../org/apache/spark/sql/internal/SQLConf.scala | 13 +++++++++++++
.../sql/execution/python/ArrowPythonRunner.scala | 1 +
.../sql/execution/python/ArrowPythonUDTFRunner.scala | 1 +
.../python/CoGroupedArrowPythonRunner.scala | 1 +
.../sql/execution/python/PythonPlannerRunner.scala | 1 +
.../spark/sql/execution/python/PythonUDFRunner.scala | 1 +
.../ApplyInPandasWithStatePythonRunner.scala | 1 +
.../python/streaming/PythonForeachWriter.scala | 1 +
12 files changed, 68 insertions(+), 5 deletions(-)
diff --git a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
index 1524ff455ec7..896ab6e3b19d 100644
--- a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
+++ b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
@@ -213,6 +213,7 @@ private[spark] abstract class BasePythonRunner[IN, OUT](
conf.get(PYTHON_DAEMON_KILL_WORKER_ON_FLUSH_FAILURE)
protected val hideTraceback: Boolean = false
protected val simplifiedTraceback: Boolean = false
+ protected val tracebackWithLocals: Boolean = false
protected def runnerConf: Map[String, String] = Map.empty
protected def evalConf: Map[String, String] = Map.empty
@@ -302,6 +303,9 @@ private[spark] abstract class BasePythonRunner[IN, OUT](
if (simplifiedTraceback) {
envVars.put("SPARK_SIMPLIFIED_TRACEBACK", "1")
}
+ if (tracebackWithLocals) {
+ envVars.put("SPARK_TRACEBACK_WITH_LOCALS", "1")
+ }
// SPARK-30299 this could be wrong with standalone mode when executor
// cores might not be correct because it defaults to all cores on the box.
val execCores =
execCoresProp.map(_.toInt).getOrElse(conf.get(EXECUTOR_CORES))
diff --git a/python/pyspark/sql/tests/test_udf.py
b/python/pyspark/sql/tests/test_udf.py
index ce7a42469e18..5701fb6dcda9 100644
--- a/python/pyspark/sql/tests/test_udf.py
+++ b/python/pyspark/sql/tests/test_udf.py
@@ -1232,6 +1232,20 @@ class BaseUDFTestsMixin:
with self.assertRaisesRegex(PythonException, "StopIteration"):
self.spark.range(10).select(test_udf(col("id"))).show()
+ def test_udf_traceback_with_locals(self):
+ with
self.sql_conf({"spark.sql.execution.pyspark.udf.tracebackWithLocals.enabled":
True}):
+
+ @udf("int")
+ def test_udf(a):
+ local_marker = a + 1
+ if local_marker:
+ raise ValueError("boom")
+ return local_marker
+
+ # The captured locals should include the local variable and its
value.
+ with self.assertRaisesRegex(PythonException, "local_marker = 1"):
+ self.spark.range(1).select(test_udf(col("id"))).collect()
+
def test_python_udf_segfault(self):
with
self.sql_conf({"spark.sql.execution.pyspark.udf.faulthandler.enabled": True}):
with self.assertRaisesRegex(Exception, "Segmentation fault"):
diff --git a/python/pyspark/tests/test_util.py
b/python/pyspark/tests/test_util.py
index c837c8dbdb97..686e946773b0 100644
--- a/python/pyspark/tests/test_util.py
+++ b/python/pyspark/tests/test_util.py
@@ -176,7 +176,9 @@ class HandleWorkerExceptionTests(unittest.TestCase):
from pyspark.util import handle_worker_exception
try:
- raise ValueError("test_message")
+ local_marker = "marker_value_42"
+ if local_marker:
+ raise ValueError("test_message")
except Exception as e:
with io.BytesIO() as stream:
handle_worker_exception(e, stream, hide_traceback)
@@ -206,6 +208,22 @@ class HandleWorkerExceptionTests(unittest.TestCase):
self.assertIn(self.exception_bytes, result)
self.assertNotIn(self.traceback_bytes, result)
+ @patch.dict(os.environ, {"SPARK_TRACEBACK_WITH_LOCALS": "1"})
+ def test_env_traceback_with_locals(self):
+ result = self.run_handle_worker_exception()
+ self.assertIn(self.exception_bytes, result)
+ self.assertIn(self.traceback_bytes, result)
+ # The local variable's value should be captured in the traceback.
+ self.assertIn(b"marker_value_42", result)
+
+ @patch.dict(os.environ, {"SPARK_TRACEBACK_WITH_LOCALS": ""})
+ def test_env_no_traceback_with_locals(self):
+ result = self.run_handle_worker_exception()
+ self.assertIn(self.exception_bytes, result)
+ self.assertIn(self.traceback_bytes, result)
+ # Without the environment variable, locals must not be captured.
+ self.assertNotIn(b"marker_value_42", result)
+
if __name__ == "__main__":
from pyspark.testing import main
diff --git a/python/pyspark/util.py b/python/pyspark/util.py
index 4d3c64e6a1f5..e66f6283ead1 100644
--- a/python/pyspark/util.py
+++ b/python/pyspark/util.py
@@ -527,12 +527,19 @@ def handle_worker_exception(
def format_exception() -> str:
if hide_traceback:
return "".join(traceback.format_exception_only(type(e), e))
+ tb = sys.exc_info()[-1]
if os.environ.get("SPARK_SIMPLIFIED_TRACEBACK", False):
- tb = try_simplify_traceback(sys.exc_info()[-1]) # type:
ignore[arg-type]
- if tb is not None:
+ simplified_tb = try_simplify_traceback(tb) # type:
ignore[arg-type]
+ if simplified_tb is not None:
+ tb = simplified_tb
e.__cause__ = None
- return "".join(traceback.format_exception(type(e), e, tb))
- return traceback.format_exc()
+ # We only set SPARK_TRACEBACK_WITH_LOCALS=1 for now. This equivalent
to a
+ # check for the existence of the environment variable.
+ capture_locals = bool(os.environ.get("SPARK_TRACEBACK_WITH_LOCALS",
False))
+ te = traceback.TracebackException(
+ type(e), e, tb, compact=True, capture_locals=capture_locals
+ )
+ return "".join(te.format())
try:
exc_info = format_exception()
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
index 4ea181330acb..abe8a9a59227 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
@@ -4802,6 +4802,17 @@ object SQLConf {
// show full stacktrace in tests but hide in production by default.
.createWithDefault(!Utils.isTesting)
+ val PYSPARK_TRACEBACK_WITH_LOCALS =
+ buildConf("spark.sql.execution.pyspark.udf.tracebackWithLocals.enabled")
+ .doc(
+ "When true, include the local variables in the traceback from Python
UDFs. " +
+ "Note that this config will print the value of every local variable in
the call stack, " +
+ "including sensitive data like passwords or API keys. Please use this
config with caution.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .booleanConf
+ .createWithDefault(false)
+
val PYSPARK_ARROW_VALIDATE_SCHEMA =
buildConf("spark.sql.execution.arrow.pyspark.validateSchema.enabled")
.doc(
@@ -8432,6 +8443,8 @@ class SQLConf extends Serializable with Logging with
SqlApiConf {
def pysparkSimplifiedTraceback: Boolean =
getConf(PYSPARK_SIMPLIFIED_TRACEBACK)
+ def pysparkTracebackWithLocals: Boolean =
getConf(PYSPARK_TRACEBACK_WITH_LOCALS)
+
def pysparkArrowValidateSchema: Boolean =
getConf(PYSPARK_ARROW_VALIDATE_SCHEMA)
def pandasGroupedMapAssignColumnsByName: Boolean =
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala
index df9b0e241b74..75b8465e2607 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala
@@ -67,6 +67,7 @@ abstract class BaseArrowPythonRunner[IN, OUT <: AnyRef](
override val hideTraceback: Boolean = SQLConf.get.pysparkHideTraceback
override val simplifiedTraceback: Boolean =
SQLConf.get.pysparkSimplifiedTraceback
+ override val tracebackWithLocals: Boolean =
SQLConf.get.pysparkTracebackWithLocals
override val bufferSize: Int = SQLConf.get.pandasUDFBufferSize
require(
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonUDTFRunner.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonUDTFRunner.scala
index a3c9f2f45926..59e78ae81e3b 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonUDTFRunner.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonUDTFRunner.scala
@@ -93,6 +93,7 @@ class ArrowPythonUDTFRunner(
override val hideTraceback: Boolean = SQLConf.get.pysparkHideTraceback
override val simplifiedTraceback: Boolean =
SQLConf.get.pysparkSimplifiedTraceback
+ override val tracebackWithLocals: Boolean =
SQLConf.get.pysparkTracebackWithLocals
override val bufferSize: Int = SQLConf.get.pandasUDFBufferSize
require(
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/CoGroupedArrowPythonRunner.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/CoGroupedArrowPythonRunner.scala
index be36b368719a..3bf22a84613e 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/CoGroupedArrowPythonRunner.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/CoGroupedArrowPythonRunner.scala
@@ -78,6 +78,7 @@ class CoGroupedArrowPythonRunner(
override val hideTraceback: Boolean = SQLConf.get.pysparkHideTraceback
override val simplifiedTraceback: Boolean =
SQLConf.get.pysparkSimplifiedTraceback
+ override val tracebackWithLocals: Boolean =
SQLConf.get.pysparkTracebackWithLocals
private val maxRecordsPerBatch: Int = {
val v = SQLConf.get.arrowMaxRecordsPerBatch
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonPlannerRunner.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonPlannerRunner.scala
index 5311f34f41c0..80887dd469d6 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonPlannerRunner.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonPlannerRunner.scala
@@ -63,6 +63,7 @@ abstract class PythonPlannerRunner[T](func: PythonFunction)
extends Logging {
val killWorkerOnFlushFailure: Boolean =
SQLConf.get.pythonUDFDaemonKillWorkerOnFlushFailure
val hideTraceback: Boolean = SQLConf.get.pysparkHideTraceback
val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+ val tracebackWithLocals: Boolean = SQLConf.get.pysparkTracebackWithLocals
val workerMemoryMb = SQLConf.get.pythonPlannerExecMemory
val jobArtifactUUID = JobArtifactSet.getCurrentJobArtifactState.map(_.uuid)
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonUDFRunner.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonUDFRunner.scala
index e59b942ab1b4..b2eaec1674b7 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonUDFRunner.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonUDFRunner.scala
@@ -60,6 +60,7 @@ abstract class BasePythonUDFRunner(
override val hideTraceback: Boolean = SQLConf.get.pysparkHideTraceback
override val simplifiedTraceback: Boolean =
SQLConf.get.pysparkSimplifiedTraceback
+ override val tracebackWithLocals: Boolean =
SQLConf.get.pysparkTracebackWithLocals
override val faultHandlerEnabled: Boolean =
SQLConf.get.pythonUDFWorkerFaulthandlerEnabled
override val idleTimeoutSeconds: Long =
SQLConf.get.pythonUDFWorkerIdleTimeoutSeconds
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/ApplyInPandasWithStatePythonRunner.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/ApplyInPandasWithStatePythonRunner.scala
index 786c5fc408dc..70d23ff89c35 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/ApplyInPandasWithStatePythonRunner.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/ApplyInPandasWithStatePythonRunner.scala
@@ -90,6 +90,7 @@ class ApplyInPandasWithStatePythonRunner(
override val hideTraceback: Boolean = sqlConf.pysparkHideTraceback
override val simplifiedTraceback: Boolean =
sqlConf.pysparkSimplifiedTraceback
+ override val tracebackWithLocals: Boolean =
sqlConf.pysparkTracebackWithLocals
override protected val largeVarTypes: Boolean = sqlConf.arrowUseLargeVarTypes
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/PythonForeachWriter.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/PythonForeachWriter.scala
index cc7745210a4d..ebcd49931da4 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/PythonForeachWriter.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/PythonForeachWriter.scala
@@ -111,6 +111,7 @@ class PythonForeachWriter(func: PythonFunction, schema:
StructType)
override val hideTraceback: Boolean = SQLConf.get.pysparkHideTraceback
override val simplifiedTraceback: Boolean =
SQLConf.get.pysparkSimplifiedTraceback
+ override val tracebackWithLocals: Boolean =
SQLConf.get.pysparkTracebackWithLocals
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
