[SPARK-13697][PySpark]Fix the missing module name of TransformFunctionSerializer.loads

[zsxwing]

What changes were proposed in this pull request?

Set the function's module name to __main__ if it's missing in TransformFunctionSerializer.loads.

How was this patch tested?

Manually test in the shell.

Before this patch:

>>> from pyspark.streaming import StreamingContext
>>> from pyspark.streaming.util import TransformFunction
>>> ssc = StreamingContext(sc, 1)
>>> func = TransformFunction(sc, lambda x: x, sc.serializer)
>>> func.rdd_wrapper(lambda x: x)
TransformFunction(<function <lambda> at 0x106ac8b18>)
>>> bytes = bytearray(ssc._transformerSerializer.serializer.dumps((func.func, func.rdd_wrap_func, func.deserializers))) 
>>> func2 = ssc._transformerSerializer.loads(bytes)
>>> print(func2.func.__module__)
None
>>> print(func2.rdd_wrap_func.__module__)
None
>>> 

After this patch:

>>> from pyspark.streaming import StreamingContext
>>> from pyspark.streaming.util import TransformFunction
>>> ssc = StreamingContext(sc, 1)
>>> func = TransformFunction(sc, lambda x: x, sc.serializer)
>>> func.rdd_wrapper(lambda x: x)
TransformFunction(<function <lambda> at 0x108bf1b90>)
>>> bytes = bytearray(ssc._transformerSerializer.serializer.dumps((func.func, func.rdd_wrap_func, func.deserializers))) 
>>> func2 = ssc._transformerSerializer.loads(bytes)
>>> print(func2.func.__module__)
__main__
>>> print(func2.rdd_wrap_func.__module__)
__main__
>>> 

[zsxwing]

cc @davies

[davies]

Could you fix this in CloudPicker?

[zsxwing]

It just calls pickle.loads(obj). So we cannot get f and wrap_func there.

[davies]

checkout save_function_tuple() and _fill_function in cloudpickle.py

[davies]

diff --git a/python/pyspark/cloudpickle.py b/python/pyspark/cloudpickle.py
index 95b3abc..e56e22a 100644
--- a/python/pyspark/cloudpickle.py
+++ b/python/pyspark/cloudpickle.py
@@ -241,6 +241,7 @@ class CloudPickler(Pickler):
         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

[davies]

Please add a test case in SerializationTestCase

[SparkQA]

Test build #52510 has finished for PR 11535 at commit c4e35ff.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[SparkQA]

Test build #52511 has finished for PR 11535 at commit 60c3068.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[zsxwing]

Updated and added a unit test. Thanks, @davies

[SparkQA]

Test build #52524 has finished for PR 11535 at commit 5f269ed.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[SparkQA]

Test build #52525 has finished for PR 11535 at commit c98042b.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[davies]

LGTM, merging this into master, thanks!

[davies]

Also merged into 1.4, 1.5, 1.6 branches.