Skip to content

Commit be649fd

Browse files
committed
remove map because we only need append
1 parent eccefcc commit be649fd

File tree

1 file changed

+6
-7
lines changed

1 file changed

+6
-7
lines changed

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

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,6 @@ import org.apache.spark.mllib.linalg.{Vector, Vectors}
3333
import org.apache.spark.rdd._
3434
import org.apache.spark.util.Utils
3535
import org.apache.spark.util.random.XORShiftRandom
36-
import org.apache.spark.util.collection.PrimitiveKeyOpenHashMap
3736

3837
/**
3938
* Entry in vocabulary
@@ -348,21 +347,21 @@ class Word2Vec extends Serializable with Logging {
348347
}
349348
val syn0Local = model._1
350349
val syn1Local = model._2
351-
val synOut = new PrimitiveKeyOpenHashMap[Int, Array[Float]](vocabSize * 2)
350+
val synOut = mutable.ListBuffer.empty[(Int, Array[Float])]
352351
var index = 0
353352
while(index < vocabSize) {
354353
if (syn0Modify(index) != 0) {
355-
synOut.update(index, syn0Local.slice(index * vectorSize, (index + 1) * vectorSize))
354+
synOut += ((index, syn0Local.slice(index * vectorSize, (index + 1) * vectorSize)))
356355
}
357356
if (syn1Modify(index) != 0) {
358-
synOut.update(index + vocabSize,
359-
syn1Local.slice(index * vectorSize, (index + 1) * vectorSize))
357+
synOut += ((index + vocabSize,
358+
syn1Local.slice(index * vectorSize, (index + 1) * vectorSize)))
360359
}
361360
index += 1
362361
}
363-
Iterator(synOut)
362+
synOut.toIterator
364363
}
365-
val synAgg = partial.flatMap(x => x).reduceByKey { case (v1, v2) =>
364+
val synAgg = partial.reduceByKey { case (v1, v2) =>
366365
blas.saxpy(vectorSize, 1.0f, v2, 1, v1, 1)
367366
v1
368367
}.collect()

0 commit comments

Comments
 (0)