address comments

marmbrus · Apr 6, 2015 · c6a2f8d · c6a2f8d
1 parent 9e2df24
commit c6a2f8d
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/python/pyspark/shuffle.py b/python/pyspark/shuffle.py
@@ -277,6 +277,9 @@ def mergeValues(self, iterator):
                 else:
                     batch *= 1.5
 
+        if get_used_memory() >= limit:
+            self._spill()
+
     def _partition(self, key):
         """ Return the partition for key """
         return hash((key, self._seed)) % self.partitions
@@ -310,6 +313,9 @@ def mergeCombiners(self, iterator, limit=None):
                 else:
                     batch *= 1.5
 
+        if get_used_memory() >= limit:
+            self._spill()
+
     def _spill(self):
         """
         dump already partitioned data into disks.