Prototype

Author: itholic

python/pyspark/errors/utils.py

@@ -16,11 +16,16 @@
 #
 
 import re
-from typing import Dict, Match
+import functools
+import inspect
+from typing import Any, Callable, Dict, Match, TypeVar, Type
 
 from pyspark.errors.error_classes import ERROR_CLASSES_MAP
 
 
+T = TypeVar("T")
+
+
 class ErrorClassesReader:
     """
     A reader to load error information from error_classes.py.
@@ -119,3 +124,62 @@ def get_message_template(self, error_class: str) -> str:
             message_template = main_message_template + " " + sub_message_template
 
         return message_template
+
+
+def _capture_call_site(func_name: str) -> None:
+    """
+    Capture the call site information including file name, line number, and function name.
+
+    This function updates the thread-local storage from server side (PySparkCurrentOrigin)
+    with the current call site information when a PySpark API function is called.
+
+    Parameters
+    ----------
+    func_name : str
+        The name of the PySpark API function being captured.
+
+    Notes
+    -----
+    The call site information is used to enhance error messages with the exact location
+    in the user code that led to the error.
+    """
+    from pyspark.sql.session import SparkSession
+
+    spark = SparkSession._getActiveSessionOrCreate()
+    assert spark._jvm is not None
+
+    stack = inspect.stack()
+    frame_info = stack[-1]
+    function = func_name
+    filename = frame_info.filename
+    lineno = frame_info.lineno
+    call_site = f'"{function}" was called from\n{filename}:{lineno}'
+
+    pyspark_origin = spark._jvm.org.apache.spark.sql.catalyst.trees.PySparkCurrentOrigin
+    pyspark_origin.set(call_site)
+
+
+def with_origin(func: Callable[..., Any]) -> Callable[..., Any]:
+    """
+    A decorator to capture and provide the call site information to the server side
+    when PySpark API functions are invoked.
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        # Update call site when the function is called
+        _capture_call_site(func.__name__)
+
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+def with_origin_to_class(cls: Type[T]) -> Type[T]:
+    """
+    Decorate all methods of a class with `with_origin` to capture call site information.
+    """
+    for name, method in cls.__dict__.items():
+        if callable(method) and name != "__init__":
+            setattr(cls, name, with_origin(method))
+    return cls

python/pyspark/sql/column.py

@@ -35,6 +35,7 @@
 
 from pyspark.context import SparkContext
 from pyspark.errors import PySparkAttributeError, PySparkTypeError, PySparkValueError
+from pyspark.errors.utils import with_origin_to_class
 from pyspark.sql.types import DataType
 from pyspark.sql.utils import get_active_spark_context
 
@@ -177,6 +178,7 @@ def _(
         return Column(njc)
 
     _.__doc__ = doc
+    _.__name__ = name
     return _
 
 
@@ -195,6 +197,7 @@ def _(self: "Column", other: Union["LiteralType", "DecimalLiteral"]) -> "Column"
     return _
 
 
+@with_origin_to_class
 class Column:
 
     """

sql/api/src/main/scala/org/apache/spark/sql/catalyst/trees/QueryContexts.scala

@@ -134,7 +134,8 @@ case class SQLQueryContext(
   override def callSite: String = throw SparkUnsupportedOperationException()
 }
 
-case class DataFrameQueryContext(stackTrace: Seq[StackTraceElement]) extends QueryContext {
+case class DataFrameQueryContext(stackTrace: Seq[StackTraceElement], pysparkCallSite: String)
+  extends QueryContext {
   override val contextType = QueryContextType.DataFrame
 
   override def objectType: String = throw SparkUnsupportedOperationException()
@@ -165,6 +166,12 @@ case class DataFrameQueryContext(stackTrace: Seq[StackTraceElement]) extends Que
     builder ++= " was called from\n"
     builder ++= callSite
     builder += '\n'
+
+    if (pysparkCallSite.nonEmpty) {
+      builder ++= "\n== PySpark call site ==\n"
+      builder ++= pysparkCallSite
+      builder += '\n'
+    }
     builder.result()
   }
 }

sql/api/src/main/scala/org/apache/spark/sql/catalyst/trees/origin.scala

@@ -35,7 +35,8 @@ case class Origin(
     stackTrace: Option[Array[StackTraceElement]] = None) {
 
   lazy val context: QueryContext = if (stackTrace.isDefined) {
-    DataFrameQueryContext(stackTrace.get.toImmutableArraySeq)
+    val pysparkCallSite = PySparkCurrentOrigin.get()
+    DataFrameQueryContext(stackTrace.get.toImmutableArraySeq, pysparkCallSite)
   } else {
     SQLQueryContext(
       line, startPosition, startIndex, stopIndex, sqlText, objectType, objectName)
@@ -84,3 +85,19 @@ object CurrentOrigin {
     ret
   }
 }
+
+/**
+ * Provides detailed call site information on PySpark.
+ * This information is generated in PySpark in the form of a String.
+ */
+object PySparkCurrentOrigin {
+  private val pysparkCallSite = new ThreadLocal[String]() {
+    override def initialValue(): String = ""
+  }
+
+  def set(value: String): Unit = pysparkCallSite.set(value)
+
+  def get(): String = pysparkCallSite.get()
+
+  def clear(): Unit = pysparkCallSite.remove()
+}