You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ru...@apache.org on 2026/10/13 02:48:02 UTC
(spark) branch master updated: [SPARK-59521][ML] Store word vectors directly in Word2VecModel
This is an automated email from the ASF dual-hosted git repository.
zhengruifeng pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 1be3a3aae398 [SPARK-59521][ML] Store word vectors directly in Word2VecModel
1be3a3aae398 is described below
commit 1be3a3aae398d4c31cdd4d477f2310af1be5d442
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Tue Oct 13 10:47:42 2026 +0800
[SPARK-59521][ML] Store word vectors directly in Word2VecModel
### What changes were proposed in this pull request?
This PR changes `org.apache.spark.ml.feature.Word2VecModel` to store its word index and
flattened word vectors directly instead of wrapping an
`org.apache.spark.mllib.feature.Word2VecModel`.
It also moves the synonym-search implementation into the ML `Word2VecModel` companion object,
reuses it from both the ML and MLlib models, and updates model fitting, copying, size estimation,
transformation, and persistence to use the directly stored data.
### Why are the changes needed?
The ML model only needs the word index and vector array from the legacy MLlib model.
Storing those values directly removes the unnecessary model wrapper and keeps the ML model's
implementation and synonym-search logic self-contained.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
The following focused suites pass with 17 tests:
```
build/sbt 'mllib/testOnly org.apache.spark.ml.feature.Word2VecSuite org.apache.spark.mllib.feature.Word2VecSuite'
```
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #58811 from zhengruifeng/ml-word2vec-store-vectors-dev-7.
Authored-by: Ruifeng Zheng <ru...@apache.org>
Signed-off-by: Ruifeng Zheng <ru...@foxmail.com>
---
.../org/apache/spark/ml/feature/Word2Vec.scala | 121 +++++++++++++++++----
.../org/apache/spark/mllib/feature/Word2Vec.scala | 59 +---------
.../apache/spark/ml/feature/Word2VecSuite.scala | 10 +-
3 files changed, 110 insertions(+), 80 deletions(-)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala
index be75e32c355b..1f2719de8bd7 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala
@@ -19,6 +19,9 @@ package org.apache.spark.ml.feature
import java.io.{DataInputStream, DataOutputStream}
+import scala.jdk.CollectionConverters._
+
+import com.google.common.collect.{Ordering => GuavaOrdering}
import org.apache.hadoop.fs.Path
import org.apache.spark.annotation.Since
@@ -28,7 +31,7 @@ import org.apache.spark.ml.linalg.{BLAS, SQLDataTypes, Vector, Vectors}
import org.apache.spark.ml.param._
import org.apache.spark.ml.param.shared._
import org.apache.spark.ml.util._
-import org.apache.spark.mllib.feature
+import org.apache.spark.mllib.feature.{Word2Vec => OldWord2Vec, Word2VecModel => OldWord2VecModel}
import org.apache.spark.sql.{DataFrame, Dataset, SparkSession}
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types._
@@ -174,7 +177,7 @@ final class Word2Vec @Since("1.4.0") (
transformSchema(dataset.schema, logging = true)
val input =
dataset.select($(inputCol)).rdd.map(_.getSeq[String](0))
- val wordVectors = new feature.Word2Vec()
+ val oldModel: OldWord2VecModel = new OldWord2Vec()
.setLearningRate($(stepSize))
.setMinCount($(minCount))
.setNumIterations($(maxIter))
@@ -184,7 +187,8 @@ final class Word2Vec @Since("1.4.0") (
.setWindowSize($(windowSize))
.setMaxSentenceLength($(maxSentenceLength))
.fit(input)
- copyValues(new Word2VecModel(uid, wordVectors).setParent(this))
+ copyValues(new Word2VecModel(
+ uid, oldModel.wordIndex, oldModel.wordVectors).setParent(this))
}
@Since("1.4.0")
@@ -209,21 +213,33 @@ object Word2Vec extends DefaultParamsReadable[Word2Vec] {
@Since("1.4.0")
class Word2VecModel private[ml] (
@Since("1.4.0") override val uid: String,
- @transient private val wordVectors: feature.Word2VecModel)
+ private val wordIndex: Map[String, Int],
+ private val wordVectors: Array[Float])
extends Model[Word2VecModel] with Word2VecBase with MLWritable {
import Word2VecModel._
// For ml connect only
- private[ml] def this() = this("", null)
+ private[ml] def this() = this("", null, null)
+
+ private def wordVectorSize: Int = wordVectors.length / wordIndex.size
+
+ @transient private lazy val wordVecInvNorms: Array[Float] = {
+ Array.tabulate(wordIndex.size) { i =>
+ val norm = BLAS.nativeBLAS.snrm2(wordVectorSize, wordVectors, i * wordVectorSize, 1)
+ if (norm != 0) 1 / norm else 0.0F
+ }
+ }
private[spark] override def estimatedSize: Long = {
var size = estimateMatadataSize
- if (wordVectors != null) {
+ if (wordIndex != null) {
// wordIndex: Map[String, Int]
- size += SizeEstimator.estimate(wordVectors.wordIndex)
+ size += SizeEstimator.estimate(wordIndex)
+ }
+ if (wordVectors != null) {
// wordVectors: Array[Float]
- size += SizeEstimator.estimate(wordVectors.wordVectors)
+ size += SizeEstimator.estimate(wordVectors)
}
size
}
@@ -235,7 +251,12 @@ class Word2VecModel private[ml] (
@Since("1.5.0")
@transient lazy val getVectors: DataFrame = {
val spark = SparkSession.builder().getOrCreate()
- val wordVec = wordVectors.getVectors.transform((_, vec) => Vectors.dense(vec.map(_.toDouble)))
+ val wordVec = wordIndex.map { case (word, index) =>
+ val offset = index * wordVectorSize
+ val vector = Vectors.dense(
+ wordVectors.slice(offset, offset + wordVectorSize).map(_.toDouble))
+ (word, vector)
+ }
spark.createDataFrame(wordVec.toSeq).toDF("word", "vector")
}
@@ -275,7 +296,8 @@ class Word2VecModel private[ml] (
*/
@Since("2.2.0")
def findSynonymsArray(vec: Vector, num: Int): Array[(String, Double)] = {
- wordVectors.findSynonyms(vec.toArray, num, None)
+ Word2VecModel.findSynonyms(
+ wordIndex, wordVectors, wordVecInvNorms, vec.toArray, num, None)
}
/**
@@ -286,7 +308,15 @@ class Word2VecModel private[ml] (
*/
@Since("2.2.0")
def findSynonymsArray(word: String, num: Int): Array[(String, Double)] = {
- wordVectors.findSynonyms(word, num)
+ val vector = wordIndex.get(word) match {
+ case Some(index) =>
+ val offset = index * wordVectorSize
+ wordVectors.slice(offset, offset + wordVectorSize).map(_.toDouble)
+ case None =>
+ throw new IllegalStateException(s"$word not in vocabulary")
+ }
+ Word2VecModel.findSynonyms(
+ wordIndex, wordVectors, wordVecInvNorms, vector, num, Some(word))
}
/** @group setParam */
@@ -306,7 +336,7 @@ class Word2VecModel private[ml] (
val outputSchema = transformSchema(dataset.schema, logging = true)
val bcWordVectors = dataset.sparkSession.sparkContext.broadcast(
- (wordVectors.wordIndex, wordVectors.wordVectors))
+ (wordIndex, wordVectors))
val size = $(vectorSize)
val emptyVec = Vectors.sparse(size, Array.emptyIntArray, Array.emptyDoubleArray)
val transformer = udf { sentence: Seq[String] =>
@@ -346,7 +376,7 @@ class Word2VecModel private[ml] (
@Since("1.4.1")
override def copy(extra: ParamMap): Word2VecModel = {
- val copied = new Word2VecModel(uid, wordVectors)
+ val copied = new Word2VecModel(uid, wordIndex, wordVectors)
copyValues(copied, extra).setParent(parent)
}
@@ -355,7 +385,7 @@ class Word2VecModel private[ml] (
@Since("3.0.0")
override def toString: String = {
- s"Word2VecModel: uid=$uid, numWords=${wordVectors.wordIndex.size}, " +
+ s"Word2VecModel: uid=$uid, numWords=${wordIndex.size}, " +
s"vectorSize=${$(vectorSize)}"
}
}
@@ -363,6 +393,50 @@ class Word2VecModel private[ml] (
@Since("1.6.0")
object Word2VecModel extends MLReadable[Word2VecModel] {
+ private[spark] def findSynonyms(
+ wordIndex: Map[String, Int],
+ wordVectors: Array[Float],
+ wordVecInvNorms: Array[Float],
+ vector: Array[Double],
+ num: Int,
+ wordOpt: Option[String]): Array[(String, Double)] = {
+ require(num > 0, "Number of similar words should > 0")
+ val vectorSize = wordVectors.length / wordIndex.size
+
+ val floatVec = vector.map(_.toFloat)
+ val vecNorm = BLAS.nativeBLAS.snrm2(vectorSize, floatVec, 1)
+
+ if (vecNorm == 0) {
+ wordIndex.keysIterator.map((_, 0.0))
+ .filterNot(t => wordOpt.contains(t._1))
+ .take(num)
+ .toArray
+ } else {
+ // Normalize input vector before BLAS.nativeBLAS.sgemv to avoid Inf value
+ BLAS.nativeBLAS.sscal(vectorSize, 1 / vecNorm, floatVec, 0, 1)
+
+ val cosineVec = Array.ofDim[Float](wordIndex.size)
+ BLAS.nativeBLAS.sgemv("T", vectorSize, wordIndex.size, 1.0F, wordVectors, vectorSize,
+ floatVec, 1, 0.0F, cosineVec, 1)
+
+ var i = 0
+ while (i < cosineVec.length) { cosineVec(i) *= wordVecInvNorms(i); i += 1 }
+
+ val idxOrd = new GuavaOrdering[(String, Int)] {
+ override def compare(left: (String, Int), right: (String, Int)): Int = {
+ Ordering[Float].compare(cosineVec(left._2), cosineVec(right._2))
+ }
+ }
+
+ idxOrd.greatestOf(wordIndex.iterator.asJava, num + 1)
+ .iterator.asScala
+ .map { case (word, index) => (word, cosineVec(index).toDouble) }
+ .filterNot(t => wordOpt.contains(t._1))
+ .take(num)
+ .toArray
+ }
+ }
+
private[Word2VecModel] case class Data(word: String, vector: Array[Float])
private[ml] def serializeData(data: Data, dos: DataOutputStream): Unit = {
@@ -384,13 +458,15 @@ object Word2VecModel extends MLReadable[Word2VecModel] {
override protected def saveImpl(path: String): Unit = {
DefaultParamsWriter.saveMetadata(instance, path, sparkSession)
- val wordVectors = instance.wordVectors.getVectors
val dataPath = new Path(path, "data").toString
val bufferSizeInBytes = Utils.byteStringAsBytes(
sc.conf.get(KRYO_SERIALIZER_MAX_BUFFER_SIZE.key, "64m"))
val numPartitions = Word2VecModelWriter.calculateNumberOfPartitions(
- bufferSizeInBytes, instance.wordVectors.wordIndex.size, instance.getVectorSize)
- val datum = wordVectors.toArray.map { case (word, vector) => Data(word, vector) }
+ bufferSizeInBytes, instance.wordIndex.size, instance.wordVectorSize)
+ val datum = instance.wordIndex.toArray.map { case (word, index) =>
+ val offset = index * instance.wordVectorSize
+ Data(word, instance.wordVectors.slice(offset, offset + instance.wordVectorSize))
+ }
ReadWriteUtils.saveArray[Data](dataPath, datum, sparkSession, serializeData, numPartitions)
}
}
@@ -435,20 +511,21 @@ object Word2VecModel extends MLReadable[Word2VecModel] {
val dataPath = new Path(path, "data").toString
- val oldModel = if (major < 2 || (major == 2 && minor < 2)) {
+ val (wordIndex, wordVectors) = if (major < 2 || (major == 2 && minor < 2)) {
val data = spark.read.parquet(dataPath)
.select("wordIndex", "wordVectors")
.head()
val wordIndex = data.getAs[Map[String, Int]](0)
val wordVectors = data.getAs[Seq[Float]](1).toArray
- new feature.Word2VecModel(wordIndex, wordVectors)
+ (wordIndex, wordVectors)
} else {
val datum = ReadWriteUtils.loadArray[Data](dataPath, sparkSession, deserializeData)
- val wordVectorsMap = datum.map(wordVector => (wordVector.word, wordVector.vector)).toMap
- new feature.Word2VecModel(wordVectorsMap)
+ val wordIndex = datum.iterator.map(_.word).zipWithIndex.toMap
+ val wordVectors = datum.flatMap(_.vector)
+ (wordIndex, wordVectors)
}
- val model = new Word2VecModel(metadata.uid, oldModel)
+ val model = new Word2VecModel(metadata.uid, wordIndex, wordVectors)
metadata.getAndSetParams(model)
model
}
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
index ee380636c8e4..f46dbc7b96c7 100644
--- a/mllib/src/main/scala/org/apache/spark/mllib/feature/Word2Vec.scala
+++ b/mllib/src/main/scala/org/apache/spark/mllib/feature/Word2Vec.scala
@@ -22,7 +22,6 @@ import java.lang.{Iterable => JavaIterable}
import scala.collection.mutable
import scala.jdk.CollectionConverters._
-import com.google.common.collect.{Ordering => GuavaOrdering}
import org.json4s.{DefaultFormats, Formats}
import org.json4s.JsonDSL._
import org.json4s.jackson.JsonMethods._
@@ -34,6 +33,7 @@ import org.apache.spark.broadcast.Broadcast
import org.apache.spark.internal.Logging
import org.apache.spark.internal.LogKeys.{ALPHA, COUNT, NUM_TRAIN_WORD, VOCAB_SIZE}
import org.apache.spark.internal.config.Kryo.KRYO_SERIALIZER_MAX_BUFFER_SIZE
+import org.apache.spark.ml.feature.{Word2VecModel => NewWord2VecModel}
import org.apache.spark.ml.linalg.BLAS
import org.apache.spark.mllib.linalg.{Vector, Vectors}
import org.apache.spark.mllib.util.{Loader, Saveable}
@@ -566,7 +566,8 @@ class Word2VecModel private[spark] (
@Since("1.1.0")
def findSynonyms(word: String, num: Int): Array[(String, Double)] = {
val vector = transform(word)
- findSynonyms(vector.toArray, num, Some(word))
+ NewWord2VecModel.findSynonyms(
+ wordIndex, wordVectors, wordVecInvNorms, vector.toArray, num, Some(word))
}
/**
@@ -579,58 +580,8 @@ class Word2VecModel private[spark] (
*/
@Since("1.1.0")
def findSynonyms(vector: Vector, num: Int): Array[(String, Double)] = {
- findSynonyms(vector.toArray, num, None)
- }
-
- /**
- * Find synonyms of the vector representation of a word, rejecting
- * words identical to the value of wordOpt, if one is supplied.
- * @param vector vector representation of a word
- * @param num number of synonyms to find
- * @param wordOpt optionally, a word to reject from the results list
- * @return array of (word, cosineSimilarity)
- */
- private[spark] def findSynonyms(
- vector: Array[Double],
- num: Int,
- wordOpt: Option[String]): Array[(String, Double)] = {
- require(num > 0, "Number of similar words should > 0")
- val localVectorSize = vectorSize
-
- val floatVec = vector.map(_.toFloat)
- val vecNorm = BLAS.nativeBLAS.snrm2(localVectorSize, floatVec, 1)
-
- val localNumWords = numWords
- if (vecNorm == 0) {
- wordIndex.keysIterator.map((_, 0.0))
- .filterNot(t => wordOpt.contains(t._1))
- .take(num)
- .toArray
- } else {
- // Normalize input vector before BLAS.nativeBLAS.sgemv to avoid Inf value
- BLAS.nativeBLAS.sscal(localVectorSize, 1 / vecNorm, floatVec, 0, 1)
-
- val cosineVec = Array.ofDim[Float](localNumWords)
- BLAS.nativeBLAS.sgemv("T", localVectorSize, localNumWords, 1.0F, wordVectors, localVectorSize,
- floatVec, 1, 0.0F, cosineVec, 1)
-
- val localWordVecInvNorms = wordVecInvNorms
- var i = 0
- while (i < cosineVec.length) { cosineVec(i) *= localWordVecInvNorms(i); i += 1 }
-
- val idxOrd = new GuavaOrdering[(String, Int)] {
- override def compare(left: (String, Int), right: (String, Int)): Int = {
- Ordering[Float].compare(cosineVec(left._2), cosineVec(right._2))
- }
- }
-
- idxOrd.greatestOf(wordIndex.iterator.asJava, num + 1)
- .iterator.asScala
- .map { case (word, index) => (word, cosineVec(index).toDouble) }
- .filterNot(t => wordOpt.contains(t._1))
- .take(num)
- .toArray
- }
+ NewWord2VecModel.findSynonyms(
+ wordIndex, wordVectors, wordVecInvNorms, vector.toArray, num, None)
}
/**
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/Word2VecSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/Word2VecSuite.scala
index f79bdb7dfd1f..e92a7380dcb4 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/Word2VecSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/Word2VecSuite.scala
@@ -21,7 +21,6 @@ import org.apache.spark.ml.linalg.{Vector, Vectors}
import org.apache.spark.ml.param.ParamsSuite
import org.apache.spark.ml.util.{DefaultReadWriteTest, MLTest, MLTestingUtils}
import org.apache.spark.ml.util.TestingUtils._
-import org.apache.spark.mllib.feature.{Word2VecModel => OldWord2VecModel}
import org.apache.spark.sql.Row
import org.apache.spark.util.Utils
@@ -31,7 +30,7 @@ class Word2VecSuite extends MLTest with DefaultReadWriteTest {
test("params") {
ParamsSuite.checkParams(new Word2Vec)
- val model = new Word2VecModel("w2v", new OldWord2VecModel(Map("a" -> Array(0.0f))))
+ val model = new Word2VecModel("w2v", Map("a" -> 0), Array(0.0f))
ParamsSuite.checkParams(model)
}
@@ -213,8 +212,11 @@ class Word2VecSuite extends MLTest with DefaultReadWriteTest {
("taiwan", Array(0.60f, 0.50f, 0.50f, 0.50f)),
("korea", Array(0.45f, 0.60f, 0.60f, 0.60f))
)
- val oldModel = new OldWord2VecModel(word2VecMap)
- val instance = new Word2VecModel("myWord2VecModel", oldModel)
+ val wordIndex = word2VecMap.keysIterator.zipWithIndex.toMap
+ val wordVectors = wordIndex.toSeq.sortBy(_._2).flatMap { case (word, _) =>
+ word2VecMap(word)
+ }.toArray
+ val instance = new Word2VecModel("myWord2VecModel", wordIndex, wordVectors)
val newInstance = testDefaultReadWrite(instance)
assert(newInstance.getVectors.collect().sortBy(_.getString(0)) ===
instance.getVectors.collect().sortBy(_.getString(0)))
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org