diff --git a/python/pyspark/cloudpickle.py b/python/pyspark/cloudpickle.py index 95b3abc74244b..e56e22a9b920e 100644 --- a/python/pyspark/cloudpickle.py +++ b/python/pyspark/cloudpickle.py @@ -241,6 +241,7 @@ def save_function_tuple(self, func): save(f_globals) save(defaults) save(dct) + save(func.__module__) write(pickle.TUPLE) write(pickle.REDUCE) # applies _fill_function on the tuple @@ -698,13 +699,14 @@ def _genpartial(func, args, kwds): return partial(func, *args, **kwds) -def _fill_function(func, globals, defaults, dict): +def _fill_function(func, globals, defaults, dict, module): """ Fills in the rest of function data into the skeleton function object that were created via _make_skel_func(). """ func.__globals__.update(globals) func.__defaults__ = defaults func.__dict__ = dict + func.__module__ = module return func diff --git a/python/pyspark/tests.py b/python/pyspark/tests.py index 23720502a82c8..a5a83c7e38b3c 100644 --- a/python/pyspark/tests.py +++ b/python/pyspark/tests.py @@ -228,6 +228,12 @@ def test_itemgetter(self): getter2 = ser.loads(ser.dumps(getter)) self.assertEqual(getter(d), getter2(d)) + def test_function_module_name(self): + ser = CloudPickleSerializer() + func = lambda x: x + func2 = ser.loads(ser.dumps(func)) + self.assertEqual(func.__module__, func2.__module__) + def test_attrgetter(self): from operator import attrgetter ser = CloudPickleSerializer()