Merge pull request #48 from svenkreiss/group-by-key-performance

rewrite groupByKey()
svenkreiss · Jun 18, 2016 · ed358f2 · ed358f2
2 parents 6fff30f + ccf4da8
commit ed358f2
Showing 1 changed file with 5 additions and 6 deletions.
diff --git a/pysparkling/rdd.py b/pysparkling/rdd.py
@@ -677,12 +677,11 @@ def groupByKey(self, numPartitions=None):
         if numPartitions is None:
             numPartitions = self.getNumPartitions()
 
-        return self.context.parallelize((
-            (k, [gg[1] for gg in g]) for k, g in itertools.groupby(
-                sorted(self.collect(), key=itemgetter(0)),
-                lambda e: e[0],
-            )
-        ), numPartitions)
+        r = defaultdict(list)
+        for key, value in self.toLocalIterator():
+            r[key].append(value)
+
+        return self.context.parallelize(r.items(), numPartitions)
 
     def histogram(self, buckets):
         """histogram