address comments

apache · Aug 19, 2014 · 52c37c3 · 52c37c3
1 parent 65f892d
commit 52c37c3
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/classification/NaiveBayes.scala b/mllib/src/main/scala/org/apache/spark/mllib/classification/NaiveBayes.scala
@@ -98,7 +98,7 @@ class NaiveBayes private (private var lambda: Double) extends Serializable with
         case dv: DenseVector =>
           dv.values
       }
-      if (!values.forall(x => x >= 0.0)) {
+      if (!values.forall(_ >= 0.0)) {
         throw new SparkException(s"Naive Bayes requires nonnegative feature values but found $v.")
       }
     }

diff --git a/mllib/src/test/scala/org/apache/spark/mllib/classification/NaiveBayesSuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/classification/NaiveBayesSuite.scala
@@ -114,6 +114,14 @@ class NaiveBayesSuite extends FunSuite with LocalSparkContext {
     intercept[SparkException] {
       NaiveBayes.train(sc.makeRDD(sparse, 2))
     }
+    val nan = Seq(
+      LabeledPoint(1.0, Vectors.sparse(1, Array(0), Array(1.0))),
+      LabeledPoint(0.0, Vectors.sparse(1, Array(0), Array(Double.NaN))),
+      LabeledPoint(1.0, Vectors.sparse(1, Array(0), Array(1.0))),
+      LabeledPoint(1.0, Vectors.sparse(1, Array.empty, Array.empty)))
+    intercept[SparkException] {
+      NaiveBayes.train(sc.makeRDD(nan, 2))
+    }
   }
 }