zhengruifeng commented on a change in pull request #31279:
URL: https://github.com/apache/spark/pull/31279#discussion_r564311695
##########
File path:
mllib/src/main/scala/org/apache/spark/mllib/recommendation/MatrixFactorizationModel.scala
##########
@@ -295,28 +296,40 @@ object MatrixFactorizationModel extends
Loader[MatrixFactorizationModel] {
srcFeatures: RDD[(Int, Array[Double])],
dstFeatures: RDD[(Int, Array[Double])],
num: Int): RDD[(Int, Array[(Int, Double)])] = {
+ import scala.collection.JavaConverters._
val srcBlocks = blockify(srcFeatures)
val dstBlocks = blockify(dstFeatures)
- 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 i = 0
- val pq = new BoundedPriorityQueue[(Int, Double)](n)(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.ddot(rank, srcFactor, 1, dstFactor, 1)
- pq += dstId -> score
- }
- pq.foreach { case (dstId, score) =>
- output(i) = (srcId, (dstId, score))
- i += 1
+
+ val ratings = srcBlocks.cartesian(dstBlocks)
+ .mapPartitions { iter =>
+ var scores: Array[Double] = null
+ var idxOrd: GuavaOrdering[Int] = 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 (scores == null || scores.length < n) {
+ scores = Array.ofDim[Double](n)
+ idxOrd = new GuavaOrdering[Int] {
Review comment:
This looks like the one in ml.ALS, however, in ml.ALS, vector are
represented as array[Float], here use array[Double].
----------------------------------------------------------------
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:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
