Add tests and other isinstance changes

python/pyspark/mllib/linalg.py

@@ -586,7 +586,7 @@ def dot(self, other):
 
         assert len(self) == _vector_size(other), "dimension mismatch"
 
-        if type(other) == array.array:
+        if isinstance(other, array.array):
             result = 0.0
             for i, ind in enumerate(self.indices):
                 result += self.values[i] * other[ind]
@@ -641,23 +641,23 @@ def squared_distance(self, other):
         AssertionError: dimension mismatch
         """
         assert len(self) == _vector_size(other), "dimension mismatch"
-        if type(other) in (list, array.array):
-            if type(other) is np.array and other.ndim != 1:
-                raise Exception("Cannot call squared_distance with %d-dimensional array" %
-                                other.ndim)
+        if isinstance(other, list) or isinstance(other, array.array):
             result = 0.0
             j = 0   # index into our own array
-            for i in xrange(len(other)):
+            for i, val in enumerate(other):
                 if j < len(self.indices) and self.indices[j] == i:
-                    diff = self.values[j] - other[i]
+                    diff = self.values[j] - val
                     result += diff * diff
                     j += 1
                 else:
-                    result += other[i] * other[i]
+                    result += val * val
             return result
 
-        elif type(other) in (np.array, np.ndarray, DenseVector):
-            if type(other) == DenseVector:
+        elif isinstance(other, np.ndarray) or isinstance(other, DenseVector):
+            if isinstance(other, np.ndarray) and other.ndim != 1:
+                raise Exception("Cannot call squared_distance with %d-dimensional array" %
+                                other.ndim)
+            if isinstance(other, DenseVector):
                 other = other.array
             sparse_ind = np.zeros(other.size, dtype=bool)
             sparse_ind[self.indices] = True
@@ -668,7 +668,7 @@ def squared_distance(self, other):
             result += np.dot(other_ind, other_ind)
             return result
 
-        elif type(other) is SparseVector:
+        elif isinstance(other, SparseVector):
             result = 0.0
             i, j = 0, 0
             while i < len(self.indices) and j < len(other.indices):

python/pyspark/mllib/tests.py

@@ -128,17 +128,22 @@ def test_dot(self):
                      [1., 2., 3., 4.],
                      [1., 2., 3., 4.],
                      [1., 2., 3., 4.]])
+        arr = pyarray('d', [0, 1, 2, 3])
         self.assertEquals(10.0, sv.dot(dv))
         self.assertTrue(array_equal(array([3., 6., 9., 12.]), sv.dot(mat)))
         self.assertEquals(30.0, dv.dot(dv))
         self.assertTrue(array_equal(array([10., 20., 30., 40.]), dv.dot(mat)))
         self.assertEquals(30.0, lst.dot(dv))
         self.assertTrue(array_equal(array([10., 20., 30., 40.]), lst.dot(mat)))
+        self.assertTrue(7.0, sv.dot(arr))
 
     def test_squared_distance(self):
         sv = SparseVector(4, {1: 1, 3: 2})
         dv = DenseVector(array([1., 2., 3., 4.]))
         lst = DenseVector([4, 3, 2, 1])
+        lst1 = [4, 3, 2, 1]
+        arr = pyarray('d', [0, 2, 1, 3])
+        narr = array([0, 2, 1, 3])
         self.assertEquals(15.0, _squared_distance(sv, dv))
         self.assertEquals(25.0, _squared_distance(sv, lst))
         self.assertEquals(20.0, _squared_distance(dv, lst))
@@ -148,6 +153,9 @@ def test_squared_distance(self):
         self.assertEquals(0.0, _squared_distance(sv, sv))
         self.assertEquals(0.0, _squared_distance(dv, dv))
         self.assertEquals(0.0, _squared_distance(lst, lst))
+        self.assertEquals(25.0, _squared_distance(sv, lst1))
+        self.assertEquals(3.0, _squared_distance(sv, arr))
+        self.assertEquals(3.0, _squared_distance(sv, narr))
 
     def test_conversion(self):
         # numpy arrays should be automatically upcast to float64