address comments:

1. make SameKey can be iterable multiple times
2. make SameKey picklable
3. more tests
4. mapPartitions() with preservePartitions=True

Author: davies

python/pyspark/join.py

@@ -48,7 +48,7 @@ def dispatch(seq):
                 vbuf.append(v)
             elif n == 2:
                 wbuf.append(v)
-        return [(v, w) for v in vbuf for w in wbuf]
+        return ((v, w) for v in vbuf for w in wbuf)
     return _do_python_join(rdd, other, numPartitions, dispatch)
 
 
@@ -62,7 +62,7 @@ def dispatch(seq):
                 wbuf.append(v)
         if not vbuf:
             vbuf.append(None)
-        return [(v, w) for v in vbuf for w in wbuf]
+        return ((v, w) for v in vbuf for w in wbuf)
     return _do_python_join(rdd, other, numPartitions, dispatch)
 
 
@@ -76,7 +76,7 @@ def dispatch(seq):
                 wbuf.append(v)
         if not wbuf:
             wbuf.append(None)
-        return [(v, w) for v in vbuf for w in wbuf]
+        return ((v, w) for v in vbuf for w in wbuf)
     return _do_python_join(rdd, other, numPartitions, dispatch)
 
 
@@ -88,8 +88,9 @@ def make_mapper(i):
     rdd_len = len(vrdds)
 
     def dispatch(seq):
-        bufs = [[] for i in range(rdd_len)]
-        for (n, v) in seq:
+        bufs = [[] for _ in range(rdd_len)]
+        for n, v in seq:
             bufs[n].append(v)
-        return tuple(map(ResultIterable, bufs))
+        return tuple(ResultIterable(vs) for vs in bufs)
+
     return union_vrdds.groupByKey(numPartitions).mapValues(dispatch)

python/pyspark/rdd.py

@@ -1619,7 +1619,7 @@ def groupByKey(it):
             merger.mergeCombiners(it)
             return merger.iteritems()
 
-        return shuffled.mapPartitions(groupByKey).mapValues(ResultIterable)
+        return shuffled.mapPartitions(groupByKey, True).mapValues(ResultIterable)
 
     def flatMapValues(self, f):
         """

python/pyspark/resultiterable.py

@@ -23,7 +23,8 @@
 class ResultIterable(object):
 
     """
-    A special result iterable. This is used because the standard iterator can not be pickled
+    A special result iterable. This is used because the standard
+    iterator can not be pickled
     """
 
     def __init__(self, it):
@@ -37,6 +38,3 @@ def __len__(self):
             return len(self.it)
         except TypeError:
             return sum(1 for _ in self.it)
-
-    def __reduce__(self):
-        return (ResultIterable, (list(self.it),))

python/pyspark/shuffle.py

@@ -537,59 +537,73 @@ def __init__(self, key, value, iterator, groupBy):
         self.groupBy = groupBy
         self._file = None
         self._ser = None
-        self._index = None
 
-    def __iter__(self):
-        return self
-
-    def next(self):
-        if self._index is None:
-            # begin of iterator
-            if self._file is not None:
-                if self.values:
-                    self._spill()
-                self._file.flush()
-                self._file.seek(0)
-            self._index = 0
-
-        if self._index >= len(self.values) and self._file is not None:
-            # load next chunk of values from disk
-            self.values = next(self._ser.load_stream(self._file))
-            self._index = 0
-
-        if self._index < len(self.values):
-            value = self.values[self._index]
-            self._index += 1
-            return value
+    def __getstate__(self):
+        sum(1 for _ in self)  # try to read all the values
+        if self._file is not None:
+            f = os.fdopen(os.dup(self._file.fileno()))
+            f.seek(0)
+            bytes = f.read()
+        else:
+            bytes = ''
+        return (self.key, bytes, self.values)
+
+    def __setstate__(self, item):
+        self.key, bytes, self.values = item
+        self.iterator = iter([])
+        self.groupBy = None
+        if bytes:
+            self._open_file()
+            self._file.write(bytes)
+        else:
+            self._file = None
+            self._ser = None
 
-        key, value = next(self.iterator)
-        if key == self.key:
-            return value
+    def __iter__(self):
+        if self._file is not None:
+            self._file.flush()
+            with os.fdopen(os.dup(self._file.fileno()), 'r', 65536) as f:
+                f.seek(0)
+                for values in self._ser.load_stream(f):
+                    for v in values:
+                        yield v
+
+        for v in self.values:
+            yield v
+
+        if self.groupBy and self.groupBy.next_item is None:
+            for key, value in self.iterator:
+                if key == self.key:
+                    self.append(value)  # save it for next read
+                    yield value
+                else:
+                    self.groupBy.next_item = (key, value)
+                    break
 
-        # push them back into groupBy
-        self.groupBy.next_item = (key, value)
-        raise StopIteration
+    def __len__(self):
+        return sum(1 for _ in self)
 
     def append(self, value):
-        if self._index is not None:
-            raise ValueError("Can not append value while iterating")
-
         self.values.append(value)
         # dump them into disk if the key is huge
         if len(self.values) >= 10240:
             self._spill()
 
+    def _open_file(self):
+        dirs = _get_local_dirs("objects")
+        d = dirs[id(self) % len(dirs)]
+        if not os.path.exists(d):
+            os.makedirs(d)
+        p = os.path.join(d, str(id))
+        self._file = open(p, "w+", 65536)
+        self._ser = CompressedSerializer(PickleSerializer())
+        os.unlink(p)
+
     def _spill(self):
         """ dump the values into disk """
         global MemoryBytesSpilled, DiskBytesSpilled
         if self._file is None:
-            dirs = _get_local_dirs("objects")
-            d = dirs[id(self) % len(dirs)]
-            if not os.path.exists(d):
-                os.makedirs(d)
-            p = os.path.join(d, str(id))
-            self._file = open(p, "w+", 65536)
-            self._ser = CompressedSerializer(PickleSerializer())
+            self._open_file()
 
         used_memory = get_used_memory()
         pos = self._file.tell()
@@ -600,6 +614,19 @@ def _spill(self):
         MemoryBytesSpilled += (used_memory - get_used_memory()) << 20
 
 
+class ChainedIterable(object):
+    """
+    Pickable chained iterator
+    """
+    def __init__(self, iterators):
+        self.iterators = iterators
+
+    def __iter__(self):
+        for vs in self.iterators:
+            for v in vs:
+                yield v
+
+
 class GroupByKey(object):
     """
     group a sorted iterator into [(k1, it1), (k2, it2), ...]
@@ -719,7 +746,7 @@ def _merged_items(self, index, limit=0):
         # if the memory can not hold all the partition,
         # then use sort based merge. Because of compression,
         # the data on disks will be much smaller than needed memory
-        if (size >> 20) > self.memory_limit / 10:
+        if (size >> 20) >= self.memory_limit / 10:
             return self._sorted_items(index)
 
         self.data = {}
@@ -750,8 +777,7 @@ def load_partition(j):
             sorter = ExternalSorter(self.memory_limit, ser)
             sorted_items = sorter.sorted(itertools.chain(*disk_items),
                                          key=operator.itemgetter(0))
-
-        return ((k, itertools.chain.from_iterable(vs)) for k, vs in GroupByKey(sorted_items))
+        return ((k, ChainedIterable(vs)) for k, vs in GroupByKey(sorted_items))
 
 
 if __name__ == "__main__":

python/pyspark/tests.py

@@ -31,7 +31,7 @@
 import time
 import zipfile
 import random
-from platform import python_implementation
+import itertools
 
 if sys.version_info[:2] <= (2, 6):
     import unittest2 as unittest
@@ -122,6 +122,35 @@ def test_huge_dataset(self):
                          self.N * 10)
         m._cleanup()
 
+    def test_group_by_key(self):
+
+        def gen_data(N, step):
+            for i in range(1, N + 1, step):
+                for j in range(i * 10):
+                    yield (i, j)
+
+        def gen_gs(N, step=1):
+            return shuffle.GroupByKey(gen_data(N, step))
+
+        self.assertEqual(1, len(list(gen_gs(1))))
+        self.assertEqual(2, len(list(gen_gs(2))))
+        self.assertEqual(100, len(list(gen_gs(100))))
+        self.assertEqual(range(1, 101), [k for k, _ in gen_gs(100)])
+        self.assertTrue(all(k * 10 == len(list(vs)) for k, vs in gen_gs(100)))
+
+        for k, vs in gen_gs(5002, 100):
+            if k % 1000 == 1:
+                self.assertEqual(range(k), list(itertools.islice(vs, k)))
+                self.assertEqual(k * 10, sum(1 for _ in vs))
+                self.assertEqual(range(k * 9, k * 10), list(itertools.islice(vs, k * 9, k * 10)))
+                self.assertEqual(k * 10, sum(1 for _ in vs))
+
+        ser = PickleSerializer()
+        l = ser.loads(ser.dumps(list(gen_gs(5002, 1000))))
+        for k, vs in l:
+            self.assertEqual(k * 10, len(vs))
+            self.assertEqual(range(k * 10), list(vs))
+
 
 class TestSorter(unittest.TestCase):
     def test_in_memory_sort(self):
@@ -595,6 +624,19 @@ def test_distinct(self):
         self.assertEquals(result.getNumPartitions(), 5)
         self.assertEquals(result.count(), 3)
 
+    def test_external_group_by_key(self):
+        self.sc._conf.set("spark.python.worker.memory", "5m")
+        N = 200001
+        kv = self.sc.parallelize(range(N)).map(lambda x: (x % 3, x))
+        gkv = kv.groupByKey().cache()
+        self.assertEqual(3, gkv.count())
+        filtered = gkv.filter(lambda (k, vs): k == 1)
+        self.assertEqual(1, filtered.count())
+        self.assertEqual([(1, N/3)], filtered.mapValues(len).collect())
+        result = filtered.collect()[0][1]
+        self.assertEqual(N/3, len(result))
+        self.assertTrue(isinstance(result.it, shuffle.ChainedIterable))
+
 
 class TestSQL(PySparkTestCase):