(spark) branch master updated: [SPARK-59152][ML] Create CountVectorizer broadcast per transform

[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 065397779ce4 [SPARK-59152][ML] Create CountVectorizer broadcast per transform
065397779ce4 is described below

commit 065397779ce409996abed999e1d6b46ee98a6c34
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Wed Sep 2 08:41:59 2026 +0800

    [SPARK-59152][ML] Create CountVectorizer broadcast per transform
    
    ### What changes were proposed in this pull request?
    
    This PR creates the `CountVectorizerModel` dictionary and its broadcast for each `transform` call
    instead of caching the broadcast on the model.
    
    ### Why are the changes needed?
    
    The dictionary broadcast is transformation-only execution state. Retaining it on the model keeps the
    broadcast alive for the model's lifetime and ties later transformations to the `SparkContext` used by
    the first transformation. A transform-scoped broadcast follows the lifecycle of the returned lazy
    DataFrame. This is consistent with the other explicit inference broadcasts in
    `org.apache.spark.ml`: `Word2VecModel` broadcasts its word-vector model per transform, while
    `RandomForestRegressionModel` and `GBTRegressionModel` broadcast their models per transform and
    share each broadcast between the output UDFs created by that transform.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No.
    
    ### How was this patch tested?
    
    No tests were added because this is an internal broadcast lifecycle change and the transformation
    logic is unchanged. The patch was checked with `git diff --check` and source line-length and
    non-ASCII scans.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Generated-by: OpenAI Codex (GPT-5)
    
    Closes #58452 from zhengruifeng/ml-countvectorizer-transform-broadcast-dev-5.
    
    Authored-by: Ruifeng Zheng <ru...@apache.org>
    Signed-off-by: Ruifeng Zheng <ru...@foxmail.com>
---
 .../apache/spark/ml/feature/CountVectorizer.scala  | 25 ++++++++--------------
 1 file changed, 9 insertions(+), 16 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala
index bd1917bb9e34..35ec4afcb421 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala
@@ -22,7 +22,6 @@ import java.util.{HashMap => JHashMap}
 import org.apache.hadoop.fs.Path
 
 import org.apache.spark.annotation.Since
-import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.ml.{Estimator, Model}
 import org.apache.spark.ml.attribute.{Attribute, AttributeGroup, NumericAttribute}
 import org.apache.spark.ml.linalg.{SQLDataTypes, Vectors}
@@ -297,31 +296,25 @@ class CountVectorizerModel(
   @Since("2.0.0")
   def setBinary(value: Boolean): this.type = set(binary, value)
 
-  /** Dictionary created from [[vocabulary]] and its indices, broadcast once for [[transform()]] */
-  private var broadcastDict: Option[Broadcast[JHashMap[String, Integer]]] = None
-
   @Since("2.0.0")
   override def transform(dataset: Dataset[_]): DataFrame = {
     val outputSchema = transformSchema(dataset.schema, logging = true)
-    if (broadcastDict.isEmpty) {
-      val dict = new JHashMap[String, Integer](math.ceil(vocabulary.length / 0.75).toInt)
-      var index = 0
-      while (index < vocabulary.length) {
-        dict.put(vocabulary(index), index)
-        index += 1
-      }
-      broadcastDict = Some(dataset.sparkSession.sparkContext.broadcast(dict))
+    val dict = new JHashMap[String, Integer](math.ceil(vocabulary.length / 0.75).toInt)
+    var index = 0
+    while (index < vocabulary.length) {
+      dict.put(vocabulary(index), index)
+      index += 1
     }
-    val dictBr = broadcastDict.get
+    val bcDict = dataset.sparkSession.sparkContext.broadcast(dict)
     val localMinTF = $(minTF)
     val localBinary = $(binary)
     val vectorizer = udf { document: Seq[String] =>
-      val dict = dictBr.value
-      val dictSize = dict.size()
+      val localDict = bcDict.value
+      val dictSize = localDict.size()
       val termCounts = new OpenHashMap[Int, Int]
       var tokenCount = 0L
       document.foreach { term =>
-        val index = dict.get(term)
+        val index = localDict.get(term)
         if (index != null) {
           termCounts.changeValue(index.intValue(), 1, _ + 1)
         }


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