improve runtime performance of solution

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

@@ -38,41 +38,6 @@ private[spark] object BLAS extends Serializable with Logging {
     _f2jBLAS
   }
 
-  // equivalent to xIndices.union(yIndices).distinct.sortBy(i => i)
-  def unitedIndices(xSortedIndices: Array[Int], ySortedIndices: Array[Int]): Array[Int] = {
-    val arr = new Array[Int](xSortedIndices.length + ySortedIndices.length)
-
-    var xj = 0
-    var yj = 0
-    var j = 0
-    var previ = Int.MaxValue
-
-    def getAt(arr: Array[Int], j: Int): Int = if (j < arr.length) arr(j) else Int.MaxValue
-
-    while (xj < xSortedIndices.length || yj < ySortedIndices.length) {
-      val xi = getAt(xSortedIndices, xj)
-      val yi = getAt(ySortedIndices, yj)
-
-      val i = if (xi <= yi) {
-        xj += 1
-        xi
-      }
-      else {
-        yj += 1
-        yi
-      }
-
-      if (previ != i) {
-        arr(j) = i
-        j += 1
-      }
-
-      previ = i
-    }
-
-    arr.slice(0, j)
-  }
-
   /**
    * y += a * x
    */
@@ -153,16 +118,56 @@ private[spark] object BLAS extends Serializable with Logging {
     * y += a * x
     */
   private def axpy(a: Double, x: SparseVector, y: SparseVector): Unit = {
-    require(x.size == y.size)
+    val xSortedIndices = x.indices
+    val xValues = x.values
 
-    val xIndices = x.indices
-    val yIndices = y.indices
+    val ySortedIndices = y.indices
+    val yValues = y.values
+
+    val newIndices = new Array[Int](xSortedIndices.length + ySortedIndices.length)
+    val newValues = new Array[Double](xValues.length + yValues.length)
+
+    assert(newIndices.length == newValues.length)
+
+    var xj = 0
+    var yj = 0
+    var j = 0
+    var previ = Int.MinValue
+
+    def getAt(indices: Array[Int], j: Int): Int =
+      if (j < indices.length) indices(j) else Int.MaxValue
+
+    while (xj < xSortedIndices.length || yj < ySortedIndices.length) {
+      val xi = getAt(xSortedIndices, xj)
+      val yi = getAt(ySortedIndices, yj)
+
+      val (i, value) = if (xi <= yi) {
+        val vv = a*xValues(xj)
+        xj += 1
+        (xi, vv)
+      }
+      else {
+        val vv = yValues(yj)
+        yj += 1
+        (yi, vv)
+      }
+
+      assert(i >= previ)
+
+      if (previ != i) {
+        newIndices(j) = i
+        newValues(j) = value
+        j += 1
+      }
+      else {
+        assert(newIndices(j - 1) == i)
+        newValues(j - 1) += value
+      }
 
-    val newIndices = unitedIndices(xIndices, yIndices)
-    assert(newIndices.size >= yIndices.size)
+      previ = i
+    }
 
-    val newValues = newIndices.map(i => a*x(i) + y(i))
-    y.reassign(newIndices, newValues)
+    y.reassign(newIndices.slice(0, j), newValues.slice(0, j))
   }
 
   /** Y += a * x */

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

@@ -64,22 +64,17 @@ class BLASSuite extends SparkFunSuite {
     assert(dx ~== Vectors.dense(0.1, 0.0, -0.2) absTol 1e-15)
   }
 
-  test("unitedIndices") {
-    val indices1 = Array(0, 2, 4, 4, 7, 8, 15)
-    val indices2 = Array(4, 9, 100)
-
-    val united = unitedIndices(indices1, indices2)
-
-    assert(united.length == 8)
-
-    assert(united(0) == 0)
-    assert(united(1) == 2)
-    assert(united(2) == 4)
-    assert(united(3) == 7)
-    assert(united(4) == 8)
-    assert(united(5) == 9)
-    assert(united(6) == 15)
-    assert(united(7) == 100)
+  test("axpy(a: Double, x: SparseVector, y: SparseVector)") {
+    val x = new SparseVector(25, Array(0, 2, 4, 7, 8, 15), Array(0.0, 2.0, 4.0, 7.0, 8.0, 15.0))
+    val y = new SparseVector(25, Array(4, 9, 20), Array(4.0, 9.0, 20.0))
+
+    axpy(0.5, x, y)
+
+    val expected = Vectors.sparse(25,
+      Array(0, 2, 4, 7, 8, 9, 15, 20),
+      Array(0.0, 1.0, 6.0, 3.5, 4.0, 9.0, 7.5, 20.0))
+
+    assert(expected ~== y absTol 1e-15)
   }
 
   test("axpy") {