srowen commented on a change in pull request #31279:
URL: https://github.com/apache/spark/pull/31279#discussion_r563778222
##########
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:
I suppose this can be defined outside the flatMap, but it won't matter
much.
----------------------------------------------------------------
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]
