remove code duplicate

apache · Aug 10, 2017 · cebe600 · cebe600
1 parent 5dc5c89
commit cebe600
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 30 deletions.
diff --git a/mllib-local/src/main/scala/org/apache/spark/ml/linalg/Vectors.scala b/mllib-local/src/main/scala/org/apache/spark/ml/linalg/Vectors.scala
@@ -138,6 +138,12 @@ sealed trait Vector extends Serializable {
   @Since("2.0.0")
   def toSparse: SparseVector
 
+  /**
+   * Converts this vector to a sparse vector with all explicit zeros removed when the size is known.
+   */
+  @Since("2.3.0")
+  private[linalg] def toSparse(nnz: Int): SparseVector
+
   /**
    * Converts this vector to a dense vector.
    */
@@ -152,17 +158,7 @@ sealed trait Vector extends Serializable {
     val nnz = numNonzeros
     // A dense vector needs 8 * size + 8 bytes, while a sparse vector needs 12 * nnz + 20 bytes.
     if (1.5 * (nnz + 1.0) < size) {
-      val ii = new Array[Int](nnz)
-      val vv = new Array[Double](nnz)
-      var k = 0
-      foreachActive { (i, v) =>
-        if (v != 0) {
-          ii(k) = i
-          vv(k) = v
-          k += 1
-        }
-      }
-      new SparseVector(size, ii, vv)
+      toSparse(nnz)
     } else {
       toDense
     }
@@ -505,8 +501,9 @@ class DenseVector @Since("2.0.0") ( @Since("2.0.0") val values: Array[Double]) e
     nnz
   }
 
-  override def toSparse: SparseVector = {
-    val nnz = numNonzeros
+  override def toSparse: SparseVector = toSparse(numNonzeros)
+
+  private[linalg] override def toSparse(nnz: Int): SparseVector = {
     val ii = new Array[Int](nnz)
     val vv = new Array[Double](nnz)
     var k = 0
@@ -645,8 +642,9 @@ class SparseVector @Since("2.0.0") (
     nnz
   }
 
-  override def toSparse: SparseVector = {
-    val nnz = numNonzeros
+  override def toSparse: SparseVector = toSparse(numNonzeros)
+
+  private[linalg] override def toSparse(nnz: Int): SparseVector = {
     if (nnz == numActives) {
       this
     } else {

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
@@ -151,6 +151,12 @@ sealed trait Vector extends Serializable {
   @Since("1.4.0")
   def toSparse: SparseVector
 
+  /**
+   * Converts this vector to a sparse vector with all explicit zeros removed when the size is known.
+   */
+  @Since("2.3.0")
+  private[linalg] def toSparse(nnz: Int): SparseVector
+
   /**
    * Converts this vector to a dense vector.
    */
@@ -165,17 +171,7 @@ sealed trait Vector extends Serializable {
     val nnz = numNonzeros
     // A dense vector needs 8 * size + 8 bytes, while a sparse vector needs 12 * nnz + 20 bytes.
     if (1.5 * (nnz + 1.0) < size) {
-      val ii = new Array[Int](nnz)
-      val vv = new Array[Double](nnz)
-      var k = 0
-      foreachActive { (i, v) =>
-        if (v != 0) {
-          ii(k) = i
-          vv(k) = v
-          k += 1
-        }
-      }
-      new SparseVector(size, ii, vv)
+      toSparse(nnz)
     } else {
       toDense
     }
@@ -680,8 +676,10 @@ class DenseVector @Since("1.0.0") (
   }
 
   @Since("1.4.0")
-  override def toSparse: SparseVector = {
-    val nnz = numNonzeros
+  override def toSparse: SparseVector = toSparse(numNonzeros)
+
+  @Since("2.3.0")
+  private[linalg] override def toSparse(nnz: Int): SparseVector = {
     val ii = new Array[Int](nnz)
     val vv = new Array[Double](nnz)
     var k = 0
@@ -833,8 +831,10 @@ class SparseVector @Since("1.0.0") (
   }
 
   @Since("1.4.0")
-  override def toSparse: SparseVector = {
-    val nnz = numNonzeros
+  override def toSparse: SparseVector = toSparse(numNonzeros)
+
+  @Since("2.3.0")
+  private[linalg] override def toSparse(nnz: Int): SparseVector = {
     if (nnz == numActives) {
       this
     } else {