Add a validation for k

apache · Jul 1, 2015 · 7991e15 · 7991e15
1 parent c2df35d
commit 7991e15
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/mllib/src/main/scala/org/apache/spark/ml/clustering/KMeans.scala b/mllib/src/main/scala/org/apache/spark/ml/clustering/KMeans.scala
@@ -40,7 +40,7 @@ private[clustering] trait KMeansParams
    * Set the number of clusters to create (k). Default: 2.
    * @group param
    */
-  val k = new Param[Int](this, "k", "number of clusters to create")
+  val k = new Param[Int](this, "k", "number of clusters to create", (x: Int) => x > 1)
 
   /** @group getParam */
   def getK: Int = $(k)

diff --git a/mllib/src/test/scala/org/apache/spark/ml/clustering/KMeansSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/clustering/KMeansSuite.scala
@@ -58,6 +58,12 @@ class KMeansSuite extends SparkFunSuite with MLlibTestSparkContext {
     assert(kmeans.getEpsilon === 1e-4)
   }
 
+  test("parameters validation") {
+    intercept[IllegalArgumentException] {
+      new KMeans().setK(1)
+    }
+  }
+
   test("fit & transform") {
     val predictionColName = "kmeans_prediction"
     val kmeans = new KMeans().setK(k).setPredictionCol(predictionColName)