saving argspec in udf

python/pyspark/sql/udf.py

@@ -157,7 +157,17 @@ def _create_judf(self):
         spark = SparkSession.builder.getOrCreate()
         sc = spark.sparkContext
 
-        wrapped_func = _wrap_function(sc, fail_on_stopiteration(self.func), self.returnType)
+        func = fail_on_stopiteration(self.func)
+
+        # prevent inspect to fail
+        # e.g. inspect.getargspec(sum) raises
+        # TypeError: <built-in function sum> is not a Python function
+        try:
+            func._argspec = _get_argspec(self.func)
+        except TypeError:
+            pass
+
+        wrapped_func = _wrap_function(sc, func, self.returnType)
         jdt = spark._jsparkSession.parseDataType(self.returnType.json())
         judf = sc._jvm.org.apache.spark.sql.execution.python.UserDefinedPythonFunction(
             self._name, wrapped_func, jdt, self.evalType, self.deterministic)

python/pyspark/util.py

@@ -23,8 +23,6 @@
 
 __all__ = []
 
-WRAPPED_ARGSPEC_ATTR = '_wrapped_argspec'
-
 
 def _exception_message(excp):
     """Return the message from an exception as either a str or unicode object.  Supports both
@@ -57,8 +55,8 @@ def _get_argspec(f):
     """
     # `getargspec` is deprecated since python3.0 (incompatible with function annotations).
     # See SPARK-23569.
-    if hasattr(f, WRAPPED_ARGSPEC_ATTR):
-        argspec = getattr(f, WRAPPED_ARGSPEC_ATTR)
+    if hasattr(f, '_argspec'):
+        argspec = f._argspec
     elif sys.version_info[0] < 3:
         argspec = inspect.getargspec(f)
     else:
@@ -107,16 +105,6 @@ def wrapper(*args, **kwargs):
                 exc
             )
 
-    # prevent inspect to fail
-    # e.g. inspect.getargspec(sum) raises
-    # TypeError: <built-in function sum> is not a Python function
-    try:
-        argspec = _get_argspec(f)
-    except TypeError:
-        pass
-    else:
-        setattr(wrapper, WRAPPED_ARGSPEC_ATTR, _get_argspec(f))
-
     return wrapper