(spark) branch master updated: [SPARK-59340][ML] Reduce Word2VecModel broadcast size

[ru...@apache.org]

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 e7f758270c85 [SPARK-59340][ML] Reduce Word2VecModel broadcast size
e7f758270c85 is described below

commit e7f758270c85d6a0015ce65bca471a9806157bdb
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Wed Sep 9 09:50:14 2026 +0800

    [SPARK-59340][ML] Reduce Word2VecModel broadcast size
    
    ### What changes were proposed in this pull request?
    
    Broadcast only the word-index map and flattened word-vector array used by
    `Word2VecModel.transform`, instead of broadcasting the legacy MLlib `Word2VecModel` instance.
    
    ### Why are the changes needed?
    
    The legacy model also contains derived fields used by synonym lookup, including lazy caches that
    may already have been initialized. These fields are not needed when transforming sentences, but
    they become part of the broadcast when the whole model is serialized. Broadcasting the two lookup
    structures directly keeps the payload limited to the state used by the transform UDF.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No.
    
    ### How was this patch tested?
    
    Ran:
    
    ```
    JAVA_HOME=/usr/lib/jvm/java-17-openjdk-amd64 build/sbt \
      'mllib/testOnly org.apache.spark.ml.feature.Word2VecSuite'
    ```
    
    All 10 tests passed.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Generated-by: OpenAI Codex (GPT-5)
    
    Closes #58627 from zhengruifeng/ml-spark-59340-word2vec-broadcast-dev-7.
    
    Authored-by: Ruifeng Zheng <ru...@apache.org>
    Signed-off-by: Ruifeng Zheng <ru...@foxmail.com>
---
 mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala | 6 +++---
 1 file changed, 3 insertions(+), 3 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 ae1039fc81c0..be75e32c355b 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
@@ -305,15 +305,15 @@ class Word2VecModel private[ml] (
   override def transform(dataset: Dataset[_]): DataFrame = {
     val outputSchema = transformSchema(dataset.schema, logging = true)
 
-    val bcModel = dataset.sparkSession.sparkContext.broadcast(this.wordVectors)
+    val bcWordVectors = dataset.sparkSession.sparkContext.broadcast(
+      (wordVectors.wordIndex, wordVectors.wordVectors))
     val size = $(vectorSize)
     val emptyVec = Vectors.sparse(size, Array.emptyIntArray, Array.emptyDoubleArray)
     val transformer = udf { sentence: Seq[String] =>
       if (sentence.isEmpty) {
         emptyVec
       } else {
-        val wordIndices = bcModel.value.wordIndex
-        val wordVectors = bcModel.value.wordVectors
+        val (wordIndices, wordVectors) = bcWordVectors.value
         val array = Array.ofDim[Double](size)
         var count = 0
         sentence.foreach { word =>


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