Reviving callable objects support in UDF in PySpark

python/pyspark/sql/functions.py

@@ -2087,10 +2087,13 @@ def _wrapped(self):
         """
         Wrap this udf with a function and attach docstring from func
         """
-        @functools.wraps(self.func)
+        assignments = tuple(a for a in functools.WRAPPER_ASSIGNMENTS if a != "__name__")
+
+        @functools.wraps(self.func, assigned=assignments)
         def wrapper(*args):
             return self(*args)
 
+        wrapper.__name__ = self._name
         wrapper.func = self.func
         wrapper.returnType = self.returnType
 

python/pyspark/sql/tests.py

@@ -679,6 +679,19 @@ def f(x):
         self.assertEqual(f, f_.func)
         self.assertEqual(return_type, f_.returnType)
 
+        class F(object):
+            """Identity"""
+            def __call__(self, x):
+                return x
+
+        f = F()
+        return_type = IntegerType()
+        f_ = udf(f, return_type)
+
+        self.assertTrue(f.__doc__ in f_.__doc__)
+        self.assertEqual(f, f_.func)
+        self.assertEqual(return_type, f_.returnType)
+
     def test_basic_functions(self):
         rdd = self.sc.parallelize(['{"foo":"bar"}', '{"foo":"baz"}'])
         df = self.spark.read.json(rdd)