(spark) branch branch-4.x updated: [SPARK-59397][ML] Reduce serialized size of mllib Word2VecModel

[ru...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

zhengruifeng pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/branch-4.x by this push:
     new 29eed0a08afc [SPARK-59397][ML] Reduce serialized size of mllib Word2VecModel
29eed0a08afc is described below

commit 29eed0a08afc1e75ee52ce4d8b5fc83d1c70dd09
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Thu Sep 10 17:56:52 2026 +0800

    [SPARK-59397][ML] Reduce serialized size of mllib Word2VecModel
    
    ### What changes were proposed in this pull request?
    
    This PR reduces the serialized size of `mllib`'s `Word2VecModel` by:
    
    - computing `numWords` and `vectorSize` from the model data instead of storing them;
    - removing the derived `wordList` cache and iterating over `wordIndex` when finding synonyms;
    - making the lazily computed word-vector inverse norms transient.
    
    ### Why are the changes needed?
    
    These values are derived from `wordIndex` and `wordVectors`. Storing them in the serialized model
    duplicates data, including an array containing the full vocabulary, and increases the pressure from
    shipping models between the Spark Connect server and executors.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No.
    
    ### How was this patch tested?
    
    Existing Word2Vec tests:
    
    ```
    build/sbt 'mllib/testOnly *Word2VecSuite'
    ```
    
    The run passed 19 tests with no failures. No new tests were added because this is an internal
    representation refactor covered by the existing synonym, Java API, and model persistence tests.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Generated-by: OpenAI Codex (GPT-5)
    
    Closes #58688 from zhengruifeng/ml-w2v-transient-caches-dev-7.
    
    Authored-by: Ruifeng Zheng <ru...@apache.org>
    Signed-off-by: Ruifeng Zheng <ru...@foxmail.com>
    (cherry picked from commit 505aff3786d0cf5bc17f574468da22e2eb9503e0)
    Signed-off-by: Ruifeng Zheng <ru...@foxmail.com>
---
 .../org/apache/spark/mllib/feature/Word2Vec.scala  | 24 ++++++++--------------
 1 file changed, 9 insertions(+), 15 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
index d99dfd504cb8..ffaf108493e7 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
@@ -501,18 +501,13 @@ class Word2VecModel private[spark] (
     private[spark] val wordIndex: Map[String, Int],
     private[spark] val wordVectors: Array[Float]) extends Serializable with Saveable {
 
-  private val numWords = wordIndex.size
+  private def numWords: Int = wordIndex.size
   // vectorSize: Dimension of each word's vector.
-  private val vectorSize = wordVectors.length / numWords
-
-  // wordList: Ordered list of words obtained from wordIndex.
-  private lazy val wordList: Array[String] = {
-    wordIndex.toSeq.sortBy(_._2).iterator.map(_._1).toArray
-  }
+  private def vectorSize: Int = wordVectors.length / numWords
 
   // wordVecInvNorms: Array of length numWords, each value being the inverse of
   //                  Euclidean norm of the wordVector.
-  private lazy val wordVecInvNorms: Array[Float] = {
+  @transient private lazy val wordVecInvNorms: Array[Float] = {
     val size = vectorSize
     Array.tabulate(numWords) { i =>
       val norm = BLAS.nativeBLAS.snrm2(size, wordVectors, i * size, 1)
@@ -599,10 +594,9 @@ class Word2VecModel private[spark] (
     val floatVec = vector.map(_.toFloat)
     val vecNorm = BLAS.nativeBLAS.snrm2(localVectorSize, floatVec, 1)
 
-    val localWordList = wordList
     val localNumWords = numWords
     if (vecNorm == 0) {
-      Iterator.tabulate(num + 1)(i => (localWordList(i), 0.0))
+      wordIndex.keysIterator.map((_, 0.0))
         .filterNot(t => wordOpt.contains(t._1))
         .take(num)
         .toArray
@@ -618,15 +612,15 @@ class Word2VecModel private[spark] (
       var i = 0
       while (i < cosineVec.length) { cosineVec(i) *= localWordVecInvNorms(i); i += 1 }
 
-      val idxOrd = new GuavaOrdering[Int] {
-        override def compare(left: Int, right: Int): Int = {
-          Ordering[Float].compare(cosineVec(left), cosineVec(right))
+      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(Iterator.range(0, localNumWords).asJava, num + 1)
+      idxOrd.greatestOf(wordIndex.iterator.asJava, num + 1)
         .iterator.asScala
-        .map(i => (localWordList(i), cosineVec(i).toDouble))
+        .map { case (word, index) => (word, cosineVec(index).toDouble) }
         .filterNot(t => wordOpt.contains(t._1))
         .take(num)
         .toArray


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org