Modify how to cast the input data

markhamstra · Jul 15, 2015 · 39514ec · 39514ec
1 parent 8117e18
commit 39514ec
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/api/python/PythonMLLibAPI.scala b/mllib/src/main/scala/org/apache/spark/mllib/api/python/PythonMLLibAPI.scala
@@ -525,9 +525,11 @@ private[python] class PythonMLLibAPI extends Serializable {
     if (seed != null) algo.setSeed(seed)
 
     val documents = data.rdd.map(_.asScala.toArray).map { r =>
-      r(0).getClass.getSimpleName match {
-        case "Integer" => (r(0).asInstanceOf[java.lang.Integer].toLong, r(1).asInstanceOf[Vector])
-        case "Long" => (r(0).asInstanceOf[java.lang.Long].toLong, r(1).asInstanceOf[Vector])
+      r(0) match {
+        case i: java.lang.Integer =>
+          (r(0).asInstanceOf[java.lang.Integer].toLong, r(1).asInstanceOf[Vector])
+        case i: java.lang.Long =>
+          (r(0).asInstanceOf[java.lang.Long].toLong, r(1).asInstanceOf[Vector])
         case _ => throw new IllegalArgumentException("input values contains invalid type value.")
       }
     }