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 a85e9236517f7..bd4ff4cb7e0bb 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 @@ -192,7 +192,7 @@ 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 input = dataset.select($(inputCol)) val countingRequired = $(minDF) < 1.0 || $(maxDF) < 1.0 val maybeInputSize = if (countingRequired) { if (dataset.storageLevel == StorageLevel.NONE) { @@ -213,41 +213,38 @@ class CountVectorizer @Since("1.5.0") (@Since("1.5.0") override val uid: String) $(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 filteringRequired = isSet(minDF) || isSet(maxDF) - val maybeFilteredWordCounts = if (filteringRequired) { - allWordCounts.filter { case (_, (_, df)) => df >= minDf && df <= maxDf } + val termCounts = input + .select(explode(col($(inputCol))).as("word")) + .groupBy("word") + .count() + + val wordCounts = if (filteringRequired) { + val documentCounts = input + .select(explode(array_distinct(col($(inputCol)))).as("word")) + .groupBy("word") + .count() + .withColumnRenamed("count", "documentCount") + termCounts.as("termCounts") + .join( + documentCounts.as("documentCounts"), + col("termCounts.word") <=> col("documentCounts.word")) + .filter(col("documentCount") >= minDf && col("documentCount") <= maxDf) + .select(col("termCounts.word").as("word"), col("termCounts.count")) } else { - allWordCounts + termCounts } - 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) + .orderBy(col("count").desc, col("word").asc) + .limit(vocSize) + .select("word") + .collect() + .map(_.getString(0)) - if (input.getStorageLevel != StorageLevel.NONE) { + if (input.storageLevel != StorageLevel.NONE) { input.unpersist() } - wordCounts.unpersist() - if (vocab.isEmpty) { this.logWarning("The vocabulary size is empty. " + "If this was unexpected, you may wish to lower minDF (or) increase maxDF.") 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 295e96bcfe6a0..c007339980507 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(