cache words RDD in fit

Author: Ishiihara

mllib/src/main/scala/org/apache/spark/mllib/feature/Word2Vec.scala

@@ -248,7 +248,7 @@ class Word2Vec extends Serializable with Logging {
    */
   def fit[S <: Iterable[String]](dataset: RDD[S]): Word2VecModel = {
 
-    val words = dataset.flatMap(x => x)
+    val words = dataset.flatMap(x => x).cache()
 
     learnVocab(words)
 
@@ -281,7 +281,9 @@ class Word2Vec extends Serializable with Logging {
       }
     }
 
-    val newSentences = sentences.repartition(numPartitions).cache()
+    val newSentences = sentences.repartition(numPartitions)
+    words.unpersist()
+    newSentences.cache()
     val initRandom = new XORShiftRandom(seed)
     val syn0Global =
       Array.fill[Float](vocabSize * vectorSize)((initRandom.nextFloat() - 0.5f) / vectorSize)