update test and fix a bug in topByKey

nemccarthy · May 1, 2015 · f864f5e · f864f5e
1 parent c5e0181
commit f864f5e
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 12 deletions.
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/rdd/MLPairRDDFunctions.scala b/mllib/src/main/scala/org/apache/spark/mllib/rdd/MLPairRDDFunctions.scala
@@ -39,7 +39,7 @@ class MLPairRDDFunctions[K: ClassTag, V: ClassTag](self: RDD[(K, V)]) extends Se
    * @return an RDD that contains the top k values for each key
    */
   def topByKey(num: Int)(implicit ord: Ordering[V]): RDD[(K, Array[V])] = {
-    self.aggregateByKey(new BoundedPriorityQueue[V](num)(ord))(
+    self.aggregateByKey(new BoundedPriorityQueue[V](num)(ord.reverse))(
       seqOp = (queue, item) => {
         queue += item
         queue

diff --git a/mllib/src/main/scala/org/apache/spark/mllib/recommendation/MatrixFactorizationModel.scala b/mllib/src/main/scala/org/apache/spark/mllib/recommendation/MatrixFactorizationModel.scala
@@ -142,7 +142,7 @@ class MatrixFactorizationModel(
   }
 
     /**
-   * Recommends topK products for all users
+   * Recommends topK products for all users.
    *
    * @param num how many products to return for every user.
    * @return [(Int, Array[Rating])] objects, where every tuple contains a userID and an array of
@@ -159,7 +159,7 @@ class MatrixFactorizationModel(
 
 
   /**
-   * Recommends topK users for all products
+   * Recommends topK users for all products.
    *
    * @param num how many users to return for every product.
    * @return [(Int, Array[Rating])] objects, where every tuple contains a productID and an array
@@ -216,7 +216,7 @@ object MatrixFactorizationModel extends Loader[MatrixFactorizationModel] {
         val output = new Array[(Int, (Double, Int))](m * n)
         var k = 0
         ratings.foreachActive { (i, j, r) =>
-          output(k) = (i, (r, j))
+          output(k) = (srcIds(i), (r, dstIds(j)))
           k += 1
         }
         output.toSeq

diff --git a/.../src/test/scala/org/apache/spark/mllib/recommendation/MatrixFactorizationModelSuite.scala b/.../src/test/scala/org/apache/spark/mllib/recommendation/MatrixFactorizationModelSuite.scala
@@ -76,20 +76,20 @@ class MatrixFactorizationModelSuite extends FunSuite with MLlibTestSparkContext
   test("batch predict API recommendProductsForUsers") {
     val model = new MatrixFactorizationModel(rank, userFeatures, prodFeatures)
     val topK = 10
-    val predictedProducts = model.recommendProductsForUsers(topK).collect().toMap
+    val recommendations = model.recommendProductsForUsers(topK).collectAsMap()
 
-    assert(predictedProducts(0)(0).rating ~== 17.0 relTol 1e-14)
-    assert(predictedProducts(1)(0).rating ~== 39.0 relTol 1e-14)
+    assert(recommendations(0)(0).rating ~== 17.0 relTol 1e-14)
+    assert(recommendations(1)(0).rating ~== 39.0 relTol 1e-14)
   }
 
   test("batch predict API recommendUsersForProducts") {
     val model = new MatrixFactorizationModel(rank, userFeatures, prodFeatures)
     val topK = 10
-    val predictedUsers = model.recommendUsersForProducts(topK).collect().toMap
+    val recommendations = model.recommendUsersForProducts(topK).collectAsMap()
 
-    assert(predictedUsers(2)(0).user == 1)
-    assert(predictedUsers(2)(0).rating ~== 39.0 relTol 1e-14)
-    assert(predictedUsers(2)(1).user == 0)
-    assert(predictedUsers(2)(1).rating ~== 17.0 relTol 1e-14)
+    assert(recommendations(2)(0).user == 1)
+    assert(recommendations(2)(0).rating ~== 39.0 relTol 1e-14)
+    assert(recommendations(2)(1).user == 0)
+    assert(recommendations(2)(1).rating ~== 17.0 relTol 1e-14)
   }
 }