(spark) branch master updated: [SPARK-58911][ML] Use DataFrame aggregations in CountVectorizer fit

[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 619c0677fc98 [SPARK-58911][ML] Use DataFrame aggregations in CountVectorizer fit
619c0677fc98 is described below

commit 619c0677fc983efc36df64636d5794c45a960acb
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Fri Aug 21 08:50:59 2026 +0800

    [SPARK-58911][ML] Use DataFrame aggregations in CountVectorizer fit
    
    ### What changes were proposed in this pull request?
    
    This PR rewrites `CountVectorizer.fit` to stay in the DataFrame execution path instead of
    converting the input column to an RDD.
    
    The new implementation:
    
    - computes fractional `minDF` / `maxDF` thresholds with a scalar subquery over the input size;
    - computes word counts with `explode` and grouped aggregation;
    - computes document frequency, when DF filtering is requested, with a two-stage aggregation:
      first by `(docId, word)` and then by `word`;
    - uses DataFrame ordering and `limit` to select the vocabulary.
    
    ### Why are the changes needed?
    
    The current implementation builds per-document maps manually and combines them with `reduceByKey`.
    Keeping the fit logic in DataFrame operations lets Spark plan the aggregations natively while
    preserving the existing CountVectorizer semantics for word counts, document frequency filtering,
    and deterministic vocabulary ordering.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No.
    
    ### How was this patch tested?
    
    Added regression coverage for document frequency when a document contains repeated words.
    
    ```
    ./build/sbt "mllib/testOnly org.apache.spark.ml.feature.CountVectorizerSuite"
    ```
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Generated-by: Codex (GPT-5)
    
    Closes #57755 from zhengruifeng/count-vectorizer-dataframe-fit-dev3.
    
    Authored-by: Ruifeng Zheng <ru...@apache.org>
    Signed-off-by: Ruifeng Zheng <ru...@foxmail.com>
---
 .../apache/spark/ml/feature/CountVectorizer.scala  | 80 ++++++++--------------
 .../spark/ml/feature/CountVectorizerSuite.scala    | 17 +++++
 2 files changed, 46 insertions(+), 51 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 f7d63faeea79..0c4a42c0fd57 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
@@ -31,7 +31,6 @@ import org.apache.spark.ml.util._
 import org.apache.spark.sql.{DataFrame, Dataset}
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.types._
-import org.apache.spark.storage.StorageLevel
 import org.apache.spark.util.ArrayImplicits._
 import org.apache.spark.util.SizeEstimator
 import org.apache.spark.util.collection.{OpenHashMap, Utils}
@@ -192,61 +191,40 @@ class CountVectorizer @Since("1.5.0") (@Since("1.5.0") override val uid: String)
     }
 
     val vocSize = $(vocabSize)
-    val input = dataset.select($(inputCol)).rdd.map(_.getSeq[String](0))
-    val countingRequired = $(minDF) < 1.0 || $(maxDF) < 1.0
-    val maybeInputSize = if (countingRequired) {
-      if (dataset.storageLevel == StorageLevel.NONE) {
-        input.persist(StorageLevel.MEMORY_AND_DISK)
-      }
-      Some(input.count())
-    } else {
-      None
-    }
-    val minDf = if ($(minDF) >= 1.0) {
-      $(minDF)
-    } else {
-      $(minDF) * maybeInputSize.get
-    }
-    val maxDf = if ($(maxDF) >= 1.0) {
-      $(maxDF)
-    } else {
-      $(maxDF) * maybeInputSize.get
-    }
-    require(maxDf >= minDf, "maxDF must be >= minDF.")
-    val allWordCounts = input.flatMap { tokens =>
-      val wc = new OpenHashMap[String, Long]
-      tokens.foreach { w =>
-        wc.changeValue(w, 1L, _ + 1L)
-      }
-      wc.map { case (word, count) => (word, (count, 1)) }
-    }.reduceByKey { (wcdf1, wcdf2) =>
-      (wcdf1._1 + wcdf2._1, wcdf1._2 + wcdf2._2)
-    }
-
+    val input = dataset.select($(inputCol))
+    val inputSizeCol = input.select(count(lit(0))).scalar()
+    val minDfCol = if ($(minDF) >= 1.0) lit($(minDF)) else inputSizeCol * lit($(minDF))
+    val maxDfCol = if ($(maxDF) >= 1.0) lit($(maxDF)) else inputSizeCol * lit($(maxDF))
     val filteringRequired = isSet(minDF) || isSet(maxDF)
-    val maybeFilteredWordCounts = if (filteringRequired) {
-      allWordCounts.filter { case (_, (_, df)) => df >= minDf && df <= maxDf }
+    val wordCounts = if (filteringRequired) {
+      input
+        .select(
+          monotonically_increasing_id().as("docId"),
+          col($(inputCol)).as("doc"))
+        .select(
+          col("docId"),
+          explode(col("doc")).as("word"))
+        .groupBy("docId", "word")
+        .agg(count(lit(0)).as("wordCountInDoc"))
+        .groupBy("word")
+        .agg(
+          sum("wordCountInDoc").as("count"),
+          count(lit(0)).as("docCount"))
+        .filter(minDfCol <= col("docCount") && col("docCount") <= maxDfCol)
+        .select("word", "count")
     } else {
-      allWordCounts
+      input
+        .select(explode(col($(inputCol))).as("word"))
+        .groupBy("word")
+        .agg(count(lit(0)).as("count"))
     }
 
-    val wordCounts = maybeFilteredWordCounts
-      .map { case (word, (count, _)) => (word, count) }
-      .persist(StorageLevel.MEMORY_AND_DISK)
-
-    val fullVocabSize = wordCounts.count()
-
-    val ordering = Ordering.Tuple2(Ordering.Long, Ordering.String.reverse)
-      .on[(String, Long)] { case (word, count) => (count, word) }
-
     val vocab = wordCounts
-      .top(math.min(fullVocabSize, vocSize).toInt)(ordering)
-      .map(_._1)
-
-    if (input.getStorageLevel != StorageLevel.NONE) {
-      input.unpersist()
-    }
-    wordCounts.unpersist()
+      .orderBy(col("count").desc, col("word").asc)
+      .limit(vocSize)
+      .select("word")
+      .collect()
+      .map(_.getString(0))
 
     if (vocab.isEmpty) {
       this.logWarning("The vocabulary size is empty. " +
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/CountVectorizerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/CountVectorizerSuite.scala
index 295e96bcfe6a..c00733998050 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/CountVectorizerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/CountVectorizerSuite.scala
@@ -160,6 +160,23 @@ class CountVectorizerSuite extends MLTest with DefaultReadWriteTest {
     }
   }
 
+  test("CountVectorizer document frequency ignores duplicate tokens") {
+    val df = Seq(
+      Array("a", "a", "a", "b"),
+      Array("a", "b", "b", "b"),
+      Array("b")
+    ).toDF("words")
+
+    val cvModel = new CountVectorizer()
+      .setInputCol("words")
+      .setOutputCol("features")
+      .setMinDF(2)
+      .setMaxDF(2)
+      .fit(df)
+
+    assert(cvModel.vocabulary === Array("a"))
+  }
+
   test("CountVectorizer using both minDF and maxDF") {
     // Ignore terms with count more than 3 AND less than 2
     val df = Seq(


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