use pickle for seq string SerDe

Author: Ishiihara

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

@@ -46,7 +46,6 @@ import org.apache.spark.mllib.util.MLUtils
 import org.apache.spark.rdd.RDD
 import org.apache.spark.util.Utils
 
-
 /**
  * :: DeveloperApi ::
  * The Java stubs necessary for the Python mllib bindings.
@@ -290,13 +289,13 @@ class PythonMLLibAPI extends Serializable {
    * handle to the Java object instead of the content of the Java object.
    * Extra care needs to be taken in the Python code to ensure it gets freed on
    * exit; see the Py4J documentation.
-   * @param dataBytesJRDD Input JavaRDD
+   * @param dataJRDD Input JavaRDD
    * @return A handle to java Word2VecModel instance at python side
    */
   def trainWord2Vec(
-    dataBytesJRDD: JavaRDD[Array[Byte]]
+    dataJRDD: JavaRDD[java.util.ArrayList[String]]
     ): Word2VecModel = {
-    val data = dataBytesJRDD.rdd.map(SerDe.deserializeSeqString)
+    val data = dataJRDD.rdd.map(_.toArray(new Array[String](0)).toSeq).cache()
     val word2vec = new Word2Vec()
     val model = word2vec.fit(data)
     model
@@ -311,8 +310,8 @@ class PythonMLLibAPI extends Serializable {
   def Word2VecModelTransform(
     model: Word2VecModel,
     word: String
-    ): Array[Byte] = {
-    SerDe.serializeDoubleVector(model.transform(word))
+    ): Vector = {
+    model.transform(word)
   }
 
   /**
@@ -332,8 +331,8 @@ class PythonMLLibAPI extends Serializable {
     val similarity = Vectors.dense(result.map(_._2))
     val words = result.map(_._1)
     val ret = new java.util.LinkedList[java.lang.Object]()
-    ret.add(SerDe.serializeSeqString(words))
-    ret.add(SerDe.serializeDoubleVector(similarity))
+    ret.add(words)
+    ret.add(similarity)
     ret
   }
 
@@ -347,16 +346,15 @@ class PythonMLLibAPI extends Serializable {
    */
   def Word2VecModelSynonyms(
     model: Word2VecModel,
-    vecBytes: Array[Byte],
+    vec: Vector,
     num: Int
     ): java.util.List[java.lang.Object] = {
-    val vec = SerDe.deserializeDoubleVector(vecBytes)
     val result = model.findSynonyms(vec, num)
     val similarity = Vectors.dense(result.map(_._2))
     val words = result.map(_._1)
     val ret = new java.util.LinkedList[java.lang.Object]()
-    ret.add(SerDe.serializeSeqString(words))
-    ret.add(SerDe.serializeDoubleVector(similarity))
+    ret.add(words)
+    ret.add(similarity)
     ret
   }
 

mllib/src/test/scala/org/apache/spark/mllib/api/python/PythonMLLibAPISuite.scala

@@ -91,11 +91,4 @@ class PythonMLLibAPISuite extends FunSuite {
     assert(bytes.length / 10 < 25) //  25 bytes per rating
 
   }
-
-  test("string seq serialization") {
-    val original = Array[String]("abc", "def", "ghi")
-    val bytes = SerDe.serializeSeqString(original)
-    val ss = SerDe.deserializeSeqString(bytes)
-    assert(ss === original)
-  }
 }

python/pyspark/mllib/Word2Vec.py

@@ -19,11 +19,11 @@
 Python package for Word2Vec in MLlib.
 """
 
-from pyspark.mllib._common import \
-    _serialize_double_vector, \
-    _deserialize_double_vector, \
-    _deserialize_string_seq, \
-    _get_unmangled_string_seq_rdd
+from functools import wraps
+
+from pyspark import PickleSerializer
+
+from pyspark.mllib.linalg import _convert_to_vector
 
 __all__ = ['Word2Vec', 'Word2VecModel']
 
@@ -46,18 +46,19 @@ def __del__(self):
     def transform(self, word):
         pythonAPI = self._sc._jvm.PythonMLLibAPI()
         result = pythonAPI.Word2VecModelTransform(self._java_model, word)
-        return _deserialize_double_vector(result)
+        return PickleSerializer().loads(str(self._sc._jvm.SerDe.dumps(result)))
 
     def findSynonyms(self, x, num):
+        SerDe = self._sc._jvm.SerDe
+        ser = PickleSerializer()
         pythonAPI = self._sc._jvm.PythonMLLibAPI()
         if type(x) == str:
-            result = pythonAPI.Word2VecModelSynonyms(self._java_model, x, num)
+            jlist = pythonAPI.Word2VecModelSynonyms(self._java_model, x, num)
         else:
-            xSer = _serialize_double_vector(x)
-            result = pythonAPI.Word2VecModelSynonyms(self._java_model, xSer, num)
-        words = _deserialize_string_seq(result[0])
-        similarity = _deserialize_double_vector(result[1])
-        return zip(words, similarity)
+            bytes = bytearray(ser.dumps(_convert_to_vector(x)))
+            vec = self._sc._jvm.SerDe.loads(bytes)
+            jlist = pythonAPI.Word2VecModelSynonyms(self._java_model, vec, num)
+        return PickleSerializer().loads(str(self._sc._jvm.SerDe.dumps(jlist)))
 
 
 class Word2Vec(object):
@@ -104,8 +105,7 @@ def fit(self, data):
         :param data: Input RDD
         """
         sc = data.context
-        dataBytes = _get_unmangled_string_seq_rdd(data)
-        model = sc._jvm.PythonMLLibAPI().trainWord2Vec(dataBytes._jrdd)
+        model = sc._jvm.PythonMLLibAPI().trainWord2Vec(data._to_java_object_rdd())
         return Word2VecModel(sc, model)