Finished SPARK-4431

mllib/src/main/scala/org/apache/spark/mllib/linalg/Vectors.scala

@@ -76,6 +76,22 @@ sealed trait Vector extends Serializable {
   def copy: Vector = {
     throw new NotImplementedError(s"copy is not implemented for ${this.getClass}.")
   }
+
+  /**
+   * It will return the iterator for the active elements of dense and sparse vector as
+   * (index, value) pair. Note that foreach method can be overridden for better performance
+   * in different vector implementation.
+   *
+   * @param skippingZeros Skipping zero elements explicitly if true. It will be useful when we
+   *                      iterator through dense vector having lots of zero elements which
+   *                      we want to skip. Default is false.
+   * @return Iterator[(Int, Double)] where the first element in the tuple is the index,
+   *         and the second element is the corresponding value.
+   */
+  private[spark] def activeIterator(skippingZeros: Boolean): Iterator[(Int, Double)]
+
+  private[spark] def activeIterator: Iterator[(Int, Double)] = activeIterator(false)
+
 }
 
 /**
@@ -273,6 +289,47 @@ class DenseVector(val values: Array[Double]) extends Vector {
   override def copy: DenseVector = {
     new DenseVector(values.clone())
   }
+
+  private[spark] override def activeIterator(skippingZeros: Boolean) = new Iterator[(Int, Double)] {
+    private var i = 0
+    private val valuesSize = values.size
+
+    // If zeros are asked to be explicitly skipped, the parent `size` method is called to count
+    // the number of nonzero elements using `hasNext` and `next` methods.
+    override lazy val size: Int = if (skippingZeros) super.size else valuesSize
+
+    override def hasNext = {
+      if (skippingZeros) {
+        var found = false
+        while (!found && i < valuesSize) if (values(i) != 0.0) found = true else i += 1
+      }
+      i < valuesSize
+    }
+
+    override def next = {
+      val result = (i, values(i))
+      i += 1
+      result
+    }
+
+    override def foreach[@specialized(Unit) U](f: ((Int, Double)) => U) {
+      var i = 0
+      if (skippingZeros) {
+        while (i < valuesSize) {
+          if (values(i) != 0.0) {
+            f(i, values(i))
+          }
+          i += 1
+        }
+      } else {
+        while (i < valuesSize) {
+          f(i, values(i))
+          i += 1
+        }
+      }
+    }
+  }
+
 }
 
 /**
@@ -309,4 +366,45 @@ class SparseVector(
   }
 
   private[mllib] override def toBreeze: BV[Double] = new BSV[Double](indices, values, size)
+
+  private[spark] override def activeIterator(skippingZeros: Boolean) = new Iterator[(Int, Double)] {
+    private var i = 0
+    private val valuesSize = values.size
+
+    // If zeros are asked to be explicitly skipped, the parent `size` method is called to count
+    // the number of nonzero elements using `hasNext` and `next` methods.
+    override lazy val size: Int = if (skippingZeros) super.size else valuesSize
+
+    def hasNext = {
+      if (skippingZeros) {
+        var found = false
+        while (!found && i < valuesSize) if (values(i) != 0.0) found = true else i += 1
+      }
+      i < valuesSize
+    }
+
+    def next = {
+      val result = (indices(i), values(i))
+      i += 1
+      result
+    }
+
+    override def foreach[@specialized(Unit) U](f: ((Int, Double)) => U) {
+      var i = 0
+      if (skippingZeros) {
+        while (i < valuesSize) {
+          if (values(i) != 0.0) {
+            f(indices(i), values(i))
+          }
+          i += 1
+        }
+      } else {
+        while (i < valuesSize) {
+          f(indices(i), values(i))
+          i += 1
+        }
+      }
+    }
+  }
+
 }

mllib/src/main/scala/org/apache/spark/mllib/stat/MultivariateOnlineSummarizer.scala

@@ -53,23 +53,21 @@ class MultivariateOnlineSummarizer extends MultivariateStatisticalSummary with S
    * Adds input value to position i.
    */
   private[this] def add(i: Int, value: Double) = {
-    if (value != 0.0) {
-      if (currMax(i) < value) {
-        currMax(i) = value
-      }
-      if (currMin(i) > value) {
-        currMin(i) = value
-      }
+    if (currMax(i) < value) {
+      currMax(i) = value
+    }
+    if (currMin(i) > value) {
+      currMin(i) = value
+    }
 
-      val prevMean = currMean(i)
-      val diff = value - prevMean
-      currMean(i) = prevMean + diff / (nnz(i) + 1.0)
-      currM2n(i) += (value - currMean(i)) * diff
-      currM2(i) += value * value
-      currL1(i) += math.abs(value)
+    val prevMean = currMean(i)
+    val diff = value - prevMean
+    currMean(i) = prevMean + diff / (nnz(i) + 1.0)
+    currM2n(i) += (value - currMean(i)) * diff
+    currM2(i) += value * value
+    currL1(i) += math.abs(value)
 
-      nnz(i) += 1.0
-    }
+    nnz(i) += 1.0
   }
 
   /**
@@ -95,21 +93,8 @@ class MultivariateOnlineSummarizer extends MultivariateStatisticalSummary with S
     require(n == sample.size, s"Dimensions mismatch when adding new sample." +
       s" Expecting $n but got ${sample.size}.")
 
-    sample match {
-      case dv: DenseVector => {
-        var j = 0
-        while (j < dv.size) {
-          add(j, dv.values(j))
-          j += 1
-        }
-      }
-      case sv: SparseVector =>
-        var j = 0
-        while (j < sv.indices.size) {
-          add(sv.indices(j), sv.values(j))
-          j += 1
-        }
-      case v => throw new IllegalArgumentException("Do not support vector type " + v.getClass)
+    sample.activeIterator(true).foreach {
+      case (index, value) => add(index, value)
     }
 
     totalCnt += 1

mllib/src/test/scala/org/apache/spark/mllib/linalg/VectorsSuite.scala

@@ -173,4 +173,80 @@ class VectorsSuite extends FunSuite {
     val v = Vectors.fromBreeze(x(::, 0))
     assert(v.size === x.rows)
   }
+
+  test("activeIterator") {
+    val dv = Vectors.dense(0.0, 1.2, 3.1, 0.0)
+    val sv = Vectors.sparse(4, Seq((1, 1.2), (2, 3.1), (3, 0.0)))
+
+    // Testing if the size of iterator is correct when the zeros are explicitly skipped.
+    // The default setting will not skip any zero explicitly.
+    assert(dv.activeIterator.size === 4)
+    assert(dv.activeIterator(false).size === 4)
+    assert(dv.activeIterator(true).size === 2)
+
+    assert(sv.activeIterator.size === 3)
+    assert(sv.activeIterator(false).size === 3)
+    assert(sv.activeIterator(true).size === 2)
+
+    // Testing `hasNext` and `next`
+    val dvIter1 = dv.activeIterator(false)
+    assert(dvIter1.hasNext === true && dvIter1.next === (0, 0.0))
+    assert(dvIter1.hasNext === true && dvIter1.next === (1, 1.2))
+    assert(dvIter1.hasNext === true && dvIter1.next === (2, 3.1))
+    assert(dvIter1.hasNext === true && dvIter1.next === (3, 0.0))
+    assert(dvIter1.hasNext === false)
+
+    val dvIter2 = dv.activeIterator(true)
+    assert(dvIter2.hasNext === true && dvIter2.next === (1, 1.2))
+    assert(dvIter2.hasNext === true && dvIter2.next === (2, 3.1))
+    assert(dvIter2.hasNext === false)
+
+    val svIter1 = sv.activeIterator(false)
+    assert(svIter1.hasNext === true && svIter1.next === (1, 1.2))
+    assert(svIter1.hasNext === true && svIter1.next === (2, 3.1))
+    assert(svIter1.hasNext === true && svIter1.next === (3, 0.0))
+    assert(svIter1.hasNext === false)
+
+    val svIter2 = sv.activeIterator(true)
+    assert(svIter2.hasNext === true && svIter2.next === (1, 1.2))
+    assert(svIter2.hasNext === true && svIter2.next === (2, 3.1))
+    assert(svIter2.hasNext === false)
+
+    // Testing `foreach`
+    val dvMap1 = scala.collection.mutable.Map[Int, Double]()
+    dvIter1.foreach{
+      case (index, value) => dvMap1.put(index, value)
+    }
+    assert(dvMap1.size === 4)
+    assert(dvMap1.get(0) === Some(0.0))
+    assert(dvMap1.get(1) === Some(1.2))
+    assert(dvMap1.get(2) === Some(3.1))
+    assert(dvMap1.get(3) === Some(0.0))
+
+    val dvMap2 = scala.collection.mutable.Map[Int, Double]()
+    dvIter2.foreach{
+      case (index, value) => dvMap2.put(index, value)
+    }
+    assert(dvMap2.size === 2)
+    assert(dvMap2.get(1) === Some(1.2))
+    assert(dvMap2.get(2) === Some(3.1))
+
+    val svMap1 = scala.collection.mutable.Map[Int, Double]()
+    dvIter1.foreach{
+      case (index, value) => svMap1.put(index, value)
+    }
+    assert(svMap1.size === 4)
+    assert(svMap1.get(1) === Some(1.2))
+    assert(svMap1.get(2) === Some(3.1))
+    assert(svMap1.get(3) === Some(0.0))
+
+    val svMap2 = scala.collection.mutable.Map[Int, Double]()
+    svIter2.foreach{
+      case (index, value) => svMap2.put(index, value)
+    }
+    assert(svMap2.size === 2)
+    assert(svMap2.get(1) === Some(1.2))
+    assert(svMap2.get(2) === Some(3.1))
+
+  }
 }