add hashcode and equals

yinxusen · yinxusen · commit 166a6fffcfb9 · 2016-02-10T23:58:03.000-08:00
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
@@ -136,6 +136,18 @@ class KMeansModel private[ml] (
 
   @Since("1.6.0")
   override def write: MLWriter = new KMeansModel.KMeansModelWriter(this)
+
+  override def hashCode(): Int =
+    this.getClass.hashCode() + uid.hashCode() + clusterCenters.map(_.hashCode()).sum
+
+  override def equals(other: Any): Boolean = other match {
+    case that: KMeansModel =>
+      this.uid == that.uid &&
+        this.clusterCenters.length == that.clusterCenters.length &&
+        this.clusterCenters.zip(that.clusterCenters)
+        .foldLeft(true) { case (indicator, (v1, v2)) => indicator && (v1 == v2) }
+    case _ => false
+  }
 }
 
 @Since("1.6.0")
