[SPARK-11566] [MLLIB] [PYTHON] Refactoring GaussianMixtureModel.gaussians in Python

cc jkbradley

Author: Yu ISHIKAWA <[email protected]>

Closes #9534 from yu-iskw/SPARK-11566.

(cherry picked from commit c0e48df)
Signed-off-by: Davies Liu <[email protected]>

Author: yu-iskw

mllib/src/main/scala/org/apache/spark/mllib/api/python/GaussianMixtureModelWrapper.scala

@@ -17,14 +17,11 @@
 
 package org.apache.spark.mllib.api.python
 
-import java.util.{List => JList}
-
-import scala.collection.JavaConverters._
-import scala.collection.mutable.ArrayBuffer
+import scala.collection.JavaConverters
 
 import org.apache.spark.SparkContext
-import org.apache.spark.mllib.linalg.{Vector, Vectors, Matrix}
 import org.apache.spark.mllib.clustering.GaussianMixtureModel
+import org.apache.spark.mllib.linalg.{Vector, Vectors}
 
 /**
   * Wrapper around GaussianMixtureModel to provide helper methods in Python
@@ -36,17 +33,11 @@ private[python] class GaussianMixtureModelWrapper(model: GaussianMixtureModel) {
   /**
     * Returns gaussians as a List of Vectors and Matrices corresponding each MultivariateGaussian
     */
-  val gaussians: JList[Object] = {
-    val modelGaussians = model.gaussians
-    var i = 0
-    var mu = ArrayBuffer.empty[Vector]
-    var sigma = ArrayBuffer.empty[Matrix]
-    while (i < k) {
-      mu += modelGaussians(i).mu
-      sigma += modelGaussians(i).sigma
-      i += 1
+  val gaussians: Array[Byte] = {
+    val modelGaussians = model.gaussians.map { gaussian =>
+      Array[Any](gaussian.mu, gaussian.sigma)
     }
-    List(mu.toArray, sigma.toArray).map(_.asInstanceOf[Object]).asJava
+    SerDe.dumps(JavaConverters.seqAsJavaListConverter(modelGaussians).asJava)
   }
 
   def save(sc: SparkContext, path: String): Unit = model.save(sc, path)

python/pyspark/mllib/clustering.py

@@ -266,7 +266,7 @@ def gaussians(self):
         """
         return [
             MultivariateGaussian(gaussian[0], gaussian[1])
-            for gaussian in zip(*self.call("gaussians"))]
+            for gaussian in self.call("gaussians")]
 
     @property
     @since('1.4.0')