Prevent Unpickler from intervening each unpickling by calling close.

apache · May 3, 2019 · d7312fb · d7312fb
1 parent 5c47924
commit d7312fb
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/core/src/main/scala/org/apache/spark/api/python/SerDeUtil.scala b/core/src/main/scala/org/apache/spark/api/python/SerDeUtil.scala
@@ -186,6 +186,9 @@ private[spark] object SerDeUtil extends Logging {
       val unpickle = new Unpickler
       iter.flatMap { row =>
         val obj = unpickle.loads(row)
+        // `Opcodes.MEMOIZE` of Protocol 4 (Python 3.4+) will store objects in internal map
+        // of `Unpickler`. This map is cleared when calling `Unpickler.close()`.
+        unpickle.close()
         if (batched) {
           obj match {
             case array: Array[Any] => array.toSeq

diff --git a/python/pyspark/serializers.py b/python/pyspark/serializers.py
@@ -62,12 +62,11 @@
 if sys.version < '3':
     import cPickle as pickle
     from itertools import izip as zip, imap as map
-    pickle_protocol = 2
 else:
     import pickle
     basestring = unicode = str
     xrange = range
-    pickle_protocol = 3
+pickle_protocol = pickle.HIGHEST_PROTOCOL
 
 from pyspark import cloudpickle
 from pyspark.util import _exception_message