Use generic type to represent IDs

apache · Oct 8, 2014 · b7851cc · b7851cc
1 parent e443fee
commit b7851cc
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/evaluation/RankingMetrics.scala b/mllib/src/main/scala/org/apache/spark/mllib/evaluation/RankingMetrics.scala
@@ -17,26 +17,28 @@
 
 package org.apache.spark.mllib.evaluation
 
+import scala.reflect.ClassTag
 
 import org.apache.spark.SparkContext._
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.rdd.RDD
 
 
+
 /**
  * ::Experimental::
  * Evaluator for ranking algorithms.
  *
  * @param predictionAndLabels an RDD of (predicted ranking, ground truth set) pairs.
  */
 @Experimental
-class RankingMetrics(predictionAndLabels: RDD[(Array[Double], Array[Double])]) {
+class RankingMetrics[T: ClassTag](predictionAndLabels: RDD[(Array[T], Array[T])]) {
 
   /**
    * Returns the precsion@k for each query
    */
   lazy val precAtK: RDD[Array[Double]] = predictionAndLabels.map {case (pred, lab)=>
-    val labSet : Set[Double] = lab.toSet
+    val labSet = lab.toSet
     val n = pred.length
     val topkPrec = Array.fill[Double](n)(0.0)
     var (i, cnt) = (0, 0)
@@ -55,7 +57,7 @@ class RankingMetrics(predictionAndLabels: RDD[(Array[Double], Array[Double])]) {
    * Returns the average precision for each query
    */
   lazy val avePrec: RDD[Double] = predictionAndLabels.map {case (pred, lab) =>
-    val labSet: Set[Double] = lab.toSet
+    val labSet = lab.toSet
     var (i, cnt, precSum) = (0, 0, 0.0)
     val n = pred.length
 

diff --git a/mllib/src/test/scala/org/apache/spark/mllib/evaluation/RankingMetricsSuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/evaluation/RankingMetricsSuite.scala
@@ -25,8 +25,8 @@ class RankingMetricsSuite extends FunSuite with LocalSparkContext {
   test("Ranking metrics: map, ndcg") {
     val predictionAndLabels = sc.parallelize(
       Seq(
-        (Array[Double](1, 6, 2, 7, 8, 3, 9, 10, 4, 5), Array[Double](1, 2, 3, 4, 5)),
-        (Array[Double](4, 1, 5, 6, 2, 7, 3, 8, 9, 10), Array[Double](1, 2, 3))
+        (Array[Int](1, 6, 2, 7, 8, 3, 9, 10, 4, 5), Array[Int](1, 2, 3, 4, 5)),
+        (Array[Int](4, 1, 5, 6, 2, 7, 3, 8, 9, 10), Array[Int](1, 2, 3))
       ), 2)
     val eps: Double = 1E-5