From 7ea133add452fa7a92d2bd4b7531d370beaebd10 Mon Sep 17 00:00:00 2001
From: Yu ISHIKAWA <yuu.ishikawa@gmail.com>
Date: Tue, 30 Jun 2015 09:14:07 +0900
Subject: [PATCH] Change how to validate `initMode`

---
 .../scala/org/apache/spark/ml/clustering/KMeans.scala     | 8 +++-----
 .../org/apache/spark/ml/clustering/KMeansSuite.scala      | 3 +++
 2 files changed, 6 insertions(+), 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
index 642612c4f0191..946d3938ae200 100644
--- 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
index a1a5c10b1dadf..7bb822f20d2f4 100644
--- 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") {