Repository: spark Updated Branches: refs/heads/branch-0.9 7e4a0e1a0 -> b06001482
[SPARK-791] [PySpark] fix pickle itemgetter with cloudpickle fix the problem with pickle operator.itemgetter with multiple index. Author: Davies Liu <[email protected]> Closes #1627 from davies/itemgetter and squashes the following commits: aabd7fa [Davies Liu] fix pickle itemgetter with cloudpickle (cherry picked from commit 92ef02626e793ea853cced4cbfee316f0b748ed7) Conflicts: python/pyspark/tests.py Project: http://git-wip-us.apache.org/repos/asf/spark/repo Commit: http://git-wip-us.apache.org/repos/asf/spark/commit/b0600148 Tree: http://git-wip-us.apache.org/repos/asf/spark/tree/b0600148 Diff: http://git-wip-us.apache.org/repos/asf/spark/diff/b0600148 Branch: refs/heads/branch-0.9 Commit: b060014828e42fac258b234de4669668539718c4 Parents: 7e4a0e1 Author: Davies Liu <[email protected]> Authored: Tue Jul 29 01:02:18 2014 -0700 Committer: Josh Rosen <[email protected]> Committed: Tue Jul 29 01:15:35 2014 -0700 ---------------------------------------------------------------------- python/pyspark/cloudpickle.py | 5 +++-- python/pyspark/tests.py | 6 ++++++ 2 files changed, 9 insertions(+), 2 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/spark/blob/b0600148/python/pyspark/cloudpickle.py ---------------------------------------------------------------------- diff --git a/python/pyspark/cloudpickle.py b/python/pyspark/cloudpickle.py index eb5dbb8..ad2d35a 100644 --- a/python/pyspark/cloudpickle.py +++ b/python/pyspark/cloudpickle.py @@ -560,8 +560,9 @@ class CloudPickler(pickle.Pickler): ] - itemgetter_obj = ctypes.cast(ctypes.c_void_p(id(obj)), ctypes.POINTER(ItemGetterType)).contents - return self.save_reduce(operator.itemgetter, (itemgetter_obj.item,)) + obj = ctypes.cast(ctypes.c_void_p(id(obj)), ctypes.POINTER(ItemGetterType)).contents + return self.save_reduce(operator.itemgetter, + obj.item if obj.nitems > 1 else (obj.item,)) if PyObject_HEAD: dispatch[operator.itemgetter] = save_itemgetter http://git-wip-us.apache.org/repos/asf/spark/blob/b0600148/python/pyspark/tests.py ---------------------------------------------------------------------- diff --git a/python/pyspark/tests.py b/python/pyspark/tests.py index c741bfb..5b124d9 100644 --- a/python/pyspark/tests.py +++ b/python/pyspark/tests.py @@ -187,6 +187,12 @@ class TestRDDFunctions(PySparkTestCase): os.unlink(tempFile.name) self.assertRaises(Exception, lambda: filtered_data.count()) + def test_itemgetter(self): + rdd = self.sc.parallelize([range(10)]) + from operator import itemgetter + self.assertEqual([1], rdd.map(itemgetter(1)).collect()) + self.assertEqual([(2, 3)], rdd.map(itemgetter(2, 3)).collect()) + class TestIO(PySparkTestCase):
