remove map because we only need append

apache · Aug 19, 2014 · be649fd · be649fd · Ishiihara · Aug 19, 2014
1 parent eccefcc
commit be649fd
Showing 1 changed file with 6 additions and 7 deletions.
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/feature/Word2Vec.scala b/mllib/src/main/scala/org/apache/spark/mllib/feature/Word2Vec.scala
@@ -33,7 +33,6 @@ import org.apache.spark.mllib.linalg.{Vector, Vectors}
 import org.apache.spark.rdd._
 import org.apache.spark.util.Utils
 import org.apache.spark.util.random.XORShiftRandom
-import org.apache.spark.util.collection.PrimitiveKeyOpenHashMap
 
 /**
  *  Entry in vocabulary 
@@ -348,21 +347,21 @@ class Word2Vec extends Serializable with Logging {
         }
         val syn0Local = model._1
         val syn1Local = model._2
-        val synOut = new PrimitiveKeyOpenHashMap[Int, Array[Float]](vocabSize * 2)
+        val synOut = mutable.ListBuffer.empty[(Int, Array[Float])]
         var index = 0
         while(index < vocabSize) {
           if (syn0Modify(index) != 0) {
-            synOut.update(index, syn0Local.slice(index * vectorSize, (index + 1) * vectorSize))
+            synOut += ((index, syn0Local.slice(index * vectorSize, (index + 1) * vectorSize)))
           }
           if (syn1Modify(index) != 0) {
-            synOut.update(index + vocabSize,
-              syn1Local.slice(index * vectorSize, (index + 1) * vectorSize))
+            synOut += ((index + vocabSize,
+              syn1Local.slice(index * vectorSize, (index + 1) * vectorSize)))
           }
           index += 1
         }
-        Iterator(synOut)
+        synOut.toIterator
       }
-      val synAgg = partial.flatMap(x => x).reduceByKey { case (v1, v2) =>
+      val synAgg = partial.reduceByKey { case (v1, v2) =>
           blas.saxpy(vectorSize, 1.0f, v2, 1, v1, 1)
           v1
       }.collect()