Merge pull request #27 from lensacom/ndim_size

Added ndim and size properties to ArrayRDD; resolves #13, #14

Author: kszucs

splearn/rdd.py

@@ -327,6 +327,10 @@ def __getitem__(self, key):
         else:
             return super(ArrayRDD, self).__getitem__(key)
 
+    @property
+    def ndim(self):
+        return self._rdd.first().ndim
+
     @property
     def shape(self):
         """Returns the shape of the data."""
@@ -335,6 +339,12 @@ def shape(self):
         shape = self._rdd.map(lambda x: x.shape[0]).sum()
         return (shape,) + first[1:]
 
+    @property
+    def size(self):
+        """Returns the shape of the data.
+        """
+        return np.prod(self.shape)
+
     def toarray(self):
         """Returns the data as numpy.array from each partition."""
         return np.concatenate(self.collect())
@@ -564,12 +574,6 @@ def unblock(self):
         """
         return self._rdd.flatMap(lambda cols: zip(*cols))
 
-    @property
-    def shape(self):
-        """Returns the shape of the data.
-        """
-        return (super(DictRDD, self).get(0).shape[0], self.columns)
-
     def transform(self, fn, column=None):
         """Execute a transformation on a column or columns. Returns the modified
         DictRDD.

splearn/tests/test_rdd.py

@@ -287,16 +287,41 @@ def test_blocks_size(self):
         shapes = ArrayRDD(rdd, 66).map(lambda x: x.shape[0]).collect()
         assert_true(all(np.in1d(shapes, [66, 34])))
 
+    def test_ndim(self):
+        data = np.arange(4000)
+        shapes = [(4000),
+                  (1000, 4),
+                  (200, 10, 2),
+                  (100, 10, 2, 2)]
+        for shape in shapes:
+            reshaped = data.reshape(shape)
+            rdd = self.sc.parallelize(reshaped)
+            assert_equal(ArrayRDD(rdd).ndim, reshaped.ndim)
+
     def test_shape(self):
         data = np.arange(4000)
         shapes = [(1000, 4),
                   (200, 20),
                   (100, 40),
                   (2000, 2)]
         for shape in shapes:
-            rdd = self.sc.parallelize(data.reshape(shape))
+            reshaped = data.reshape(shape)
+            rdd = self.sc.parallelize(reshaped)
             assert_equal(ArrayRDD(rdd).shape, shape)
 
+    def test_size(self):
+        data = np.arange(4000)
+        shapes = [(1000, 4),
+                  (200, 20),
+                  (100, 40),
+                  (2000, 2)]
+        for shape in shapes:
+            reshaped = data.reshape(shape)
+            rdd = self.sc.parallelize(reshaped)
+            size = ArrayRDD(rdd).map(lambda x: x.size).sum()
+            assert_equal(size, reshaped.size)
+            assert_equal(ArrayRDD(rdd).size, reshaped.size)
+
     def test_unblocking_rdd(self):
         data = np.arange(400)
         rdd = self.sc.parallelize(data, 4)