[GitHub] [spark] srowen commented on a change in pull request #30468: [SPARK-33518][ML][WIP] Improve performance of ML ALS recommendForAll by GEMV

[GitBox <gi...@apache.org>]

srowen commented on a change in pull request #30468:
URL: https://github.com/apache/spark/pull/30468#discussion_r528775009



##########
File path: mllib/src/main/scala/org/apache/spark/ml/recommendation/TopByKeyAggregator.scala
##########
@@ -57,3 +57,57 @@ private[recommendation] class TopByKeyAggregator[K1: TypeTag, K2: TypeTag, V: Ty
 
   override def outputEncoder: Encoder[Array[(K2, V)]] = ExpressionEncoder[Array[(K2, V)]]()
 }
+
+
+/**
+ * Works on rows of the form (ScrId, DstIds, Scores).
+ * Finds the top `num` DstIds and Scores.
+ */
+private[recommendation] class TopKArrayAggregator(num: Int)
+  extends Aggregator[
+    (Int, Array[Int], Array[Float]),
+    (Array[Int], Array[Float]),
+    Array[(Int, Float)]] {
+
+  override def zero: (Array[Int], Array[Float]) = {
+    (Array.emptyIntArray, Array.emptyFloatArray)
+  }
+
+  override def reduce(
+      b: (Array[Int], Array[Float]),
+      a: (Int, Array[Int], Array[Float])): (Array[Int], Array[Float]) = {
+    merge(b, (a._2, a._3))
+  }
+
+  def merge(
+      b1: (Array[Int], Array[Float]),
+      b2: (Array[Int], Array[Float])): (Array[Int], Array[Float]) = {
+    val (ids1, scores1) = b1
+    val (ids2, socres2) = b2

Review comment:
       Nit: typo in socres2

##########
File path: mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
##########
@@ -456,34 +457,42 @@ class ALSModel private[ml] (
       num: Int,
       blockSize: Int): DataFrame = {
     import srcFactors.sparkSession.implicits._
+    import ALSModel.TopSelector
 
     val srcFactorsBlocked = blockify(srcFactors.as[(Int, Array[Float])], blockSize)
     val dstFactorsBlocked = blockify(dstFactors.as[(Int, Array[Float])], blockSize)
-    val ratings = srcFactorsBlocked.crossJoin(dstFactorsBlocked)
-      .as[(Seq[(Int, Array[Float])], Seq[(Int, Array[Float])])]
-      .flatMap { case (srcIter, dstIter) =>
-        val m = srcIter.size
-        val n = math.min(dstIter.size, num)
-        val output = new Array[(Int, Int, Float)](m * n)
-        var i = 0
-        val pq = new BoundedPriorityQueue[(Int, Float)](num)(Ordering.by(_._2))
-        srcIter.foreach { case (srcId, srcFactor) =>
-          dstIter.foreach { case (dstId, dstFactor) =>
-            // We use F2jBLAS which is faster than a call to native BLAS for vector dot product
-            val score = BLAS.f2jBLAS.sdot(rank, srcFactor, 1, dstFactor, 1)
-            pq += dstId -> score
+    val partialRecs = srcFactorsBlocked.crossJoin(dstFactorsBlocked)
+      .as[(Array[Int], Array[Float], Array[Int], Array[Float])]
+      .mapPartitions { iter =>
+        var buffer: Array[Float] = null
+        var selector: TopSelector = null
+        iter.flatMap { case (srcIds, srcMat, dstIds, dstMat) =>
+          require(srcMat.length == srcIds.length * rank)
+          require(dstMat.length == dstIds.length * rank)
+          val m = srcIds.length
+          val n = dstIds.length
+          if (buffer == null || buffer.length < n) {
+            buffer = Array.ofDim[Float](n)
+            selector = new TopSelector(buffer)
           }
-          pq.foreach { case (dstId, score) =>
-            output(i) = (srcId, dstId, score)
-            i += 1
+
+          Iterator.tabulate(m) { i =>
+            // buffer = i-th vec in srcMat * dstMat
+            BLAS.f2jBLAS.sgemv("T", rank, n, 1.0F, dstMat, 0, rank,
+              srcMat, i * rank, 1, 0.0F, buffer, 0, 1)
+            val indices = selector.selectTopKIndices(Iterator.range(0, n), num)
+            (srcIds(i), indices.map(dstIds), indices.map(buffer))
           }
-          pq.clear()
+        } ++ {

Review comment:
       I'm a little confused by this part - why null these out? the flatMap in which they are declared is done here. Maybe I misread

##########
File path: mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
##########
@@ -551,6 +563,21 @@ object ALSModel extends MLReadable[ALSModel] {
       model
     }
   }
+
+  /** select top indices based on values. */
+  private[recommendation] class TopSelector(val values: Array[Float]) {

Review comment:
       Does this need to be a class? looks like this code is called once. May be less code/indirection

##########
File path: mllib/src/main/scala/org/apache/spark/ml/recommendation/TopByKeyAggregator.scala
##########
@@ -57,3 +57,57 @@ private[recommendation] class TopByKeyAggregator[K1: TypeTag, K2: TypeTag, V: Ty
 
   override def outputEncoder: Encoder[Array[(K2, V)]] = ExpressionEncoder[Array[(K2, V)]]()
 }
+
+
+/**
+ * Works on rows of the form (ScrId, DstIds, Scores).
+ * Finds the top `num` DstIds and Scores.
+ */
+private[recommendation] class TopKArrayAggregator(num: Int)

Review comment:
       I don't love the copy-paste - are the other uses of the class above able to use your new idea?




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org