Merge pull request #320 from boazmohar/repartition_bug

base.repartition() is missing a sortByKey() to the RDD

test/test_base.py

@@ -101,4 +101,23 @@ def test_map_with_keys(eng):
     data = images.fromlist([array([[1, 1], [1, 1]]), array([[2, 2], [2, 2]])], engine=eng)
     mapped = data.map(lambda kv: kv[0] + kv[1], with_keys=True)
     assert allclose(mapped.shape, [2, 2, 2])
-    assert allclose(mapped.toarray(), [[[1, 1], [1, 1]], [[3, 3], [3, 3]]])
+    assert allclose(mapped.toarray(), [[[1, 1], [1, 1]], [[3, 3], [3, 3]]])
+
+
+def test_repartition(eng):
+    if eng is not None:
+        data = images.fromlist([array([1, 1]), array([2, 2]), array([3, 3]), array([4, 4]),
+                                array([5, 5]), array([6, 6]), array([7, 7]), array([8, 8]),
+                                array([9, 9]), array([10, 10]), array([11, 11]), array([12, 12])],
+                               engine=eng, npartitions=10)
+        assert allclose(data.first(), array([1, 1]))
+        data = data.repartition(3)
+        assert allclose(data.first(), array([1, 1]))
+
+        data = series.fromlist([array([1, 1]), array([2, 2]), array([3, 3]), array([4, 4]),
+                                array([5, 5]), array([6, 6]), array([7, 7]), array([8, 8]),
+                                array([9, 9]), array([10, 10]), array([11, 11]), array([12, 12])],
+                               engine=eng, npartitions=10)
+        assert allclose(data.first(), array([1, 1]))
+        data = data.repartition(3)
+        assert allclose(data.first(), array([1, 1]))

thunder/base.py

@@ -175,8 +175,7 @@ def repartition(self, npartitions):
             Number of partitions after repartitions.
         """
         if self.mode == 'spark':
-            self.values._rdd = self.values._rdd.repartition(npartitions)
-            return self
+            return self._constructor(self.values.repartition(npartitions)).__finalize__(self)
         else:
             notsupported(self.mode)
 

thunder/images/images.py

@@ -59,7 +59,7 @@ def first(self):
             return self.values[0]
 
         if self.mode == 'spark':
-            return self.values.tordd().values().first()
+            return self.values.first()
 
     def toblocks(self, size='150', padding=None):
         """

thunder/series/series.py

@@ -113,7 +113,7 @@ def first(self):
             return self.values[tuple(zeros(len(self.baseaxes))) + (slice(None, None),)]
 
         if self.mode == 'spark':
-            return self.values.tordd().values().first()
+            return self.values.first()
 
     def tolocal(self):
         """