[Spark-11968][ML][MLLIB]Optimize MLLIB ALS recommendForAll

mllib/src/main/scala/org/apache/spark/mllib/recommendation/MatrixFactorizationModel.scala

@@ -39,6 +39,7 @@ import org.apache.spark.mllib.util.{Loader, Saveable}
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.{Row, SparkSession}
 import org.apache.spark.storage.StorageLevel
+import org.apache.spark.util.BoundedPriorityQueue
 
 /**
  * Model representing the result of matrix factorization.
@@ -276,44 +277,53 @@ object MatrixFactorizationModel extends Loader[MatrixFactorizationModel] {
       num: Int): RDD[(Int, Array[(Int, Double)])] = {
     val srcBlocks = blockify(rank, srcFeatures)
     val dstBlocks = blockify(rank, dstFeatures)
-    val ratings = srcBlocks.cartesian(dstBlocks).flatMap {
-      case ((srcIds, srcFactors), (dstIds, dstFactors)) =>
-        val m = srcIds.length
-        val n = dstIds.length
-        val ratings = srcFactors.transpose.multiply(dstFactors)
-        val output = new Array[(Int, (Int, Double))](m * n)
-        var k = 0
-        ratings.foreachActive { (i, j, r) =>
-          output(k) = (srcIds(i), (dstIds(j), r))
-          k += 1
-        }
-        output.toSeq
+    /**
+     * Use dot to replace blas 3 gemm is the key approach to improve efficiency.
+     * By this change, we can get the topK elements of each block to reduce the GC time.
+     * Comparing with BLAS.dot, hand-written dot is high efficiency.
+     */
+    val ratings = srcBlocks.cartesian(dstBlocks).flatMap { case (srcIter, dstIter) =>
+      val m = srcIter.size
+      val n = math.min(dstIter.size, num)
+      val output = new Array[(Int, (Int, Double))](m * n)
+      var j = 0
+      srcIter.foreach { case (srcId, srcFactor) =>
+          val pq = new BoundedPriorityQueue[(Int, Double)](n)(Ordering.by(_._2))
+          dstIter.foreach { case (dstId, dstFactor) =>
+              /**
+               * The below code is equivalent to
+               * val score = blas.ddot(rank, srcFactor, 1, dstFactor, 1)
+               */
+              var score: Double = 0
+              var k = 0
+              while (k < rank) {
+                score += srcFactor(k) * dstFactor(k)
+                k += 1
+              }
+              pq += ((dstId, score))
+          }
+          val pqIter = pq.iterator
+          var i = 0
+          while (i < n) {
+            output(j + i) = (srcId, pqIter.next())
+            i += 1
+          }
+          j += n
+      }
+      output.toSeq
     }
     ratings.topByKey(num)(Ordering.by(_._2))
   }
 
   /**
-   * Blockifies features to use Level-3 BLAS.
+   * Blockifies features to improve the efficiency of cartesian product
    */
   private def blockify(
       rank: Int,
-      features: RDD[(Int, Array[Double])]): RDD[(Array[Int], DenseMatrix)] = {
+      features: RDD[(Int, Array[Double])]): RDD[Seq[(Int, Array[Double])]] = {
     val blockSize = 4096 // TODO: tune the block size
-    val blockStorage = rank * blockSize
     features.mapPartitions { iter =>
-      iter.grouped(blockSize).map { grouped =>
-        val ids = mutable.ArrayBuilder.make[Int]
-        ids.sizeHint(blockSize)
-        val factors = mutable.ArrayBuilder.make[Double]
-        factors.sizeHint(blockStorage)
-        var i = 0
-        grouped.foreach { case (id, factor) =>
-          ids += id
-          factors ++= factor
-          i += 1
-        }
-        (ids.result(), new DenseMatrix(rank, i, factors.result()))
-      }
+      iter.grouped(blockSize)
     }
   }