Change how to validate initMode

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

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
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
@@ -72,7 +72,8 @@ private[clustering] trait KMeansParams
    * (Bahmani et al., Scalable K-Means++, VLDB 2012). Default: k-means||.
    * @group param
    */
-  val initMode = new Param[String](this, "initMode", "initialization algorithm")
+  val initMode = new Param[String](this, "initMode", "initialization algorithm",
+    (value: String) => MLlibKMeans.validateInitializationMode(value))
 
   /** @group getParam */
   def getInitializationMode: String = $(initMode)
@@ -157,10 +158,7 @@ class KMeans(override val uid: String) extends Estimator[KMeansModel] with KMean
   def setK(value: Int): this.type = set(k, value)
 
   /** @group setParam */
-  def setInitializationMode(value: String): this.type = {
-    MLlibKMeans.validateInitializationMode(value)
-    set(initMode, value)
-  }
+  def setInitializationMode(value: String): this.type = set(initMode, value)
 
   /** @group setParam */
   def setInitializationSteps(value: Int): this.type = {

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
@@ -62,6 +62,9 @@ class KMeansSuite extends SparkFunSuite with MLlibTestSparkContext {
     intercept[IllegalArgumentException] {
       new KMeans().setK(1)
     }
+    intercept[IllegalArgumentException] {
+      new KMeans().setInitializationMode("no_such_a_mode")
+    }
   }
 
   test("fit & transform") {