Fixing SparseVector argmax function to ignore zero values while doing…

… the calculation.
apache · May 15, 2015 · eeda560 · eeda560
1 parent 4526acc
commit eeda560
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/linalg/Vectors.scala b/mllib/src/main/scala/org/apache/spark/mllib/linalg/Vectors.scala
@@ -728,7 +728,7 @@ class SparseVector(
       var maxValue = values(0)
 
       foreachActive { (i, v) =>
-        if(v > maxValue){
+        if(v != 0.0 && v > maxValue){
           maxIdx = i
           maxValue = v
         }

diff --git a/mllib/src/test/scala/org/apache/spark/mllib/linalg/VectorsSuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/linalg/VectorsSuite.scala
@@ -71,6 +71,10 @@ class VectorsSuite extends FunSuite {
     val vec2 = Vectors.dense(arr).asInstanceOf[DenseVector]
     val max = vec2.argmax
     assert(max === 3)
+
+    val vec3 = Vectors.dense(Array(-1.0, 0.0, -2.0, 1.0)).asInstanceOf[DenseVector]
+    val max2 = vec3.argmax
+    assert(max === 3)
   }
 
   test("sparse to array") {
@@ -87,9 +91,10 @@ class VectorsSuite extends FunSuite {
     val max = vec2.argmax
     assert(max === 3)
 
-    val vec3 = Vectors.sparse(5,Array(1,3,4),Array(1.0,.5,.7))
+    // check for case that sparse vector is created with a zero value in it by mistake
+    val vec3 = Vectors.sparse(5,Array(0, 2, 4),Array(-1.0, 0.0, -.7))
     val max2 = vec3.argmax
-    assert(max2 === 1)
+    assert(max2 === 4)
   }
 
   test("vector equals") {