Merge pull request #33 from svenkreiss/avoid-tee

avoid itertools.tee but still lazily load partition data

pysparkling/partition.py

@@ -1,7 +1,6 @@
 from __future__ import absolute_import
 
 import logging
-import itertools
 
 log = logging.getLogger(__name__)
 
@@ -12,14 +11,15 @@ def __init__(self, x, idx=None):
         self._x = x
 
     def x(self):
-        self._x, r = itertools.tee(self._x)
-        return r
+        if not isinstance(self._x, list):
+            self._x = list(self._x)
+        return self._x
 
     def hashCode(self):
         return self.index
 
     def __getstate__(self):
         return {
             'index': self.index,
-            '_x': list(self.x())
+            '_x': self.x(),
         }

pysparkling/rdd.py

@@ -47,7 +47,7 @@ class RDD(object):
     """
 
     def __init__(self, partitions, ctx):
-        self._p = partitions
+        self._p = list(partitions)
         self.context = ctx
         self._name = None
         self._rdd_id = ctx.newRddId()
@@ -64,8 +64,7 @@ def compute(self, split, task_context):
         return split.x()
 
     def partitions(self):
-        self._p, r = itertools.tee(self._p, 2)
-        return r
+        return self._p
 
     """
 
@@ -644,7 +643,7 @@ def getNumPartitions(self):
             Returns the number of partitions.
 
         """
-        return sum(1 for _ in self.partitions())
+        return len(self.partitions())
 
     def getPartitions(self):
         """
@@ -1629,7 +1628,7 @@ def takeSample(self, n):
         """
 
         rnd_entries = sorted([random.random() for _ in range(n)])
-        num_partitions = sum(1 for _ in self.partitions())
+        num_partitions = self.getNumPartitions()
 
         rnd_entries = [
             (