let spilled aggregation in Python configurable

add spark.python.worker.memory for memory used by Python worker.
Default is 512M.

Author: davies

docs/configuration.md

@@ -195,6 +195,15 @@ Apart from these, the following properties are also available, and may be useful
     Spark's dependencies and user dependencies. It is currently an experimental feature.
   </td>
 </tr>
+<tr>
+  <td><code>spark.python.worker.memory</code></td>
+  <td>512m</td>
+  <td>
+    Amount of memory to use per python worker process during aggregation, in the same
+    format as JVM memory strings (e.g. <code>512m</code>, <code>2g</code>). If the memory
+    used during aggregation go above this amount, it will spill the data into disks.
+  </td>
+</tr>
 </table>
 
 #### Shuffle Behavior

python/pyspark/rdd.py

@@ -42,7 +42,7 @@
 from pyspark.rddsampler import RDDSampler
 from pyspark.storagelevel import StorageLevel
 from pyspark.resultiterable import ResultIterable
-from pyspark.shuffle import Merger
+from pyspark.shuffle import MapMerger, ExternalHashMapMerger
 
 from py4j.java_collections import ListConverter, MapConverter
 
@@ -169,6 +169,18 @@ def _replaceRoot(self, value):
             self._sink(1)
 
 
+def _parse_memory(s):
+    """
+    >>> _parse_memory("256m")
+    256
+    >>> _parse_memory("2g")
+    2048
+    """
+    units = {'g': 1024, 'm': 1, 't': 1<<20, 'k':1.0/1024}
+    if s[-1] not in units:
+        raise ValueError("invalid format: " + s)
+    return int(float(s[:-1]) * units[s[-1].lower()])
+
 class RDD(object):
 
     """
@@ -1249,10 +1261,14 @@ def combineLocally(iterator):
         locally_combined = self.mapPartitions(combineLocally)
         shuffled = locally_combined.partitionBy(numPartitions)
 
-        executorMemory = self.ctx._jsc.sc().executorMemory()
+        serializer = self.ctx.serializer
+        spill = ((self.ctx._conf.get("spark.shuffle.spill") or 'True').lower()
+                in ('true', '1', 'yes'))
+        memory = _parse_memory(self.ctx._conf.get("spark.python.worker.memory") or "512m")
         def _mergeCombiners(iterator):
-            # TODO: workdir and serializer
-            merger = Merger(mergeCombiners, executorMemory)
+            # TODO: workdir
+            merger = ExternalHashMapMerger(mergeCombiners, memory, serializer=serializer)\
+                         if spill else MapMerger(mergeCombiners)
             merger.merge(iterator)
             return merger.iteritems()
         return shuffled.mapPartitions(_mergeCombiners)

python/pyspark/shuffle.py

@@ -45,13 +45,41 @@ def get_used_memory():
 
 
 class Merger(object):
+    """
+    merge shuffled data together by combinator
+    """
+    def merge(self, iterator):
+        raise NotImplementedError
+
+    def iteritems(self):
+        raise NotImplementedError
+
+
+class MapMerger(Merger):
+    """
+    In memory merger based on map
+    """
+    def __init__(self, combiner):
+        self.combiner = combiner
+        self.data = {}
+
+    def merge(self, iterator):
+        d, comb = self.data, self.combiner
+        for k,v in iter(iterator):
+            d[k] = comb(d[k], v) if k in d else v
+
+    def iteritems(self):
+        return self.data.iteritems()
+
+
+class ExternalHashMapMerger(Merger):
 
     """
     External merger will dump the aggregated data into disks when memory usage 
     is above the limit, then merge them together.
 
     >>> combiner = lambda x, y:x+y
-    >>> merger = Merger(combiner, 10)
+    >>> merger = ExternalHashMapMerger(combiner, 10)
     >>> N = 10000
     >>> merger.merge(zip(xrange(N), xrange(N)) * 10)
     >>> assert merger.spills > 0
@@ -63,16 +91,16 @@ class Merger(object):
     PARTITIONS = 64
     BATCH = 10000
 
-    def __init__(self, combiner, memory_limit=512, path="/tmp/pysparki/merge",
-            serializer=None, batch_size=1024, scale=1):
+    def __init__(self, combiner, memory_limit=512, path="/tmp/pyspark/merge",
+            serializer=None, scale=1):
         self.combiner = combiner
-        self.path = os.path.join(path, str(os.getpid()))
         self.memory_limit = memory_limit
-        self.serializer = serializer or BatchedSerializer(AutoSerializer(), batch_size)
+        self.path = os.path.join(path, str(os.getpid()))
+        self.serializer = serializer or BatchedSerializer(AutoSerializer(), 1024)
+        self.scale = scale
         self.data = {}
         self.pdata = []
         self.spills = 0
-        self.scale = scale
 
     @property
     def used_memory(self):
@@ -94,7 +122,7 @@ def merge(self, iterator, check=True):
                 continue
 
             c += 1
-            if c % self.BATCH == 0 and self.used_memory > self.memory_limit:
+            if c % batch == 0 and self.used_memory > self.memory_limit:
                 self._first_spill()
                 self._partitioned_merge(iterator, self.next_limit)
                 break
@@ -158,7 +186,7 @@ def _external_items(self):
             for j in range(self.spills):
                 p = os.path.join(self.path, str(j), str(i))
                 self.merge(self.serializer.load_stream(open(p)), check=False)
-                
+
                 if j > 0 and self.used_memory > hard_limit and j < self.spills - 1:
                     self.data.clear() # will read from disk again
                     for v in self._recursive_merged_items(i):
@@ -178,12 +206,12 @@ def _recursive_merged_items(self, start):
             self._spill()
 
         for i in range(start, self.PARTITIONS):
-            m = Merger(self.combiner, self.memory_limit,
+            m = ExternalHashMapMerger(self.combiner, self.memory_limit,
                     os.path.join(self.path, 'merge', str(i)),
                     self.serializer, scale=self.scale * self.PARTITIONS) 
-            m.pdata = [{} for x in range(self.PARTITIONS)]
+            m.pdata = [{} for _ in range(self.PARTITIONS)]
             limit = self.next_limit
-        
+
             for j in range(self.spills):
                 p = os.path.join(self.path, str(j), str(i))
                 m._partitioned_merge(self.serializer.load_stream(open(p)), 0)
@@ -193,7 +221,7 @@ def _recursive_merged_items(self, start):
 
             for v in m._external_items():
                 yield v
-            
+
         shutil.rmtree(self.path, True)
 
 

python/pyspark/tests.py

@@ -34,7 +34,7 @@
 from pyspark.context import SparkContext
 from pyspark.files import SparkFiles
 from pyspark.serializers import read_int
-from pyspark.shuffle import Merger
+from pyspark.shuffle import MapMerger, ExternalHashMapMerger
 
 _have_scipy = False
 try:
@@ -54,23 +54,28 @@ def setUp(self):
         self.N = 1<<18
         self.l = [i for i in xrange(self.N)]
         self.data = zip(self.l, self.l)
-        Merger.PARTITIONS = 8
-        Merger.BATCH = 1<<14
+        ExternalHashMapMerger.PARTITIONS = 8
+        ExternalHashMapMerger.BATCH = 1<<14
+
+    def test_in_memory(self):
+        m = MapMerger(lambda x,y: x+y)
+        m.merge(self.data)
+        self.assertEqual(sum(v for k,v in m.iteritems()), sum(xrange(self.N)))
 
     def test_small_dataset(self):
-        m = Merger(lambda x,y: x+y, 1000)
+        m = ExternalHashMapMerger(lambda x,y: x+y, 1000)
         m.merge(self.data)
         self.assertEqual(m.spills, 0)
         self.assertEqual(sum(v for k,v in m.iteritems()), sum(xrange(self.N)))
 
     def test_medium_dataset(self):
-        m = Merger(lambda x,y: x+y, 10)
+        m = ExternalHashMapMerger(lambda x,y: x+y, 10)
         m.merge(self.data * 3)
         self.assertTrue(m.spills >= 1)
         self.assertEqual(sum(v for k,v in m.iteritems()), sum(xrange(self.N)) * 3)
 
     def test_huge_dataset(self):
-        m = Merger(lambda x,y: x + y, 10)
+        m = ExternalHashMapMerger(lambda x,y: x + y, 10)
         m.merge(map(lambda (k,v): (k, [str(v)]), self.data) * 10)
         self.assertTrue(m.spills >= 1)
         self.assertEqual(sum(len(v) for k,v in m._recursive_merged_items(0)), self.N * 10)