Spark MLlib TF-IDF

TF-IDF is a feature extraction technique used to convert text documents into numerical vectors. In Spark MLlib, TF-IDF is commonly built by combining a term-frequency transformer such as HashingTF or CountVectorizer with the IDF estimator. In this tutorial, we shall learn what TF-IDF means, how Spark calculates TF-IDF features, and how to implement TF-IDF using Spark MLlib in Java and Python.

What is TF-IDF in Spark MLlib?

TF-IDF stands for Term Frequency – Inverse Document Frequency. It assigns a weight to each term in a document based on two ideas: how often the term appears in that document, and how rare the term is across the whole corpus.

In text mining, the input is usually a corpus containing many documents. Each document contains terms, such as words, tokens, or word-parts. TF-IDF helps represent each document as a sparse vector that can be used for text classification, clustering, similarity search, recommendation, and other machine learning tasks.

TF (Term Frequency) – In the context of term and document, TF is defined as the number of times a term appears in a document. Term and Document are independent variables and TF is dependent on these. Let us denote TF as a function of term (t) and document (d) : TF(t,d).

DF (Document Frequency) – In the context of term and all the documents in corpus, DF is defined as the number of documents that contain the term. Term and Document Corpus are independent variables and DF is dependent on these. Let us denote DF as a function of term (t) and document corpus (D) : DF(t,D).

TF alone can overemphasize frequent words. For example, words that appear in almost every document may have high term frequency, but they may not help distinguish one document from another. IDF reduces the weight of terms that are common across the corpus and gives relatively higher weight to terms that occur in fewer documents.

Note: Base of the log could be any number > 1. Spark’s ML implementation uses natural logarithm when it computes IDF internally.

Manual TF-IDF calculation for a small Spark-style corpus

Consider the following text corpus containing three documents.

document1 : Welcome to TutorialKart. There are many tutorials covering various fields of technology.

document2 : Technology has advanced a lot with the invention of semi-conductor transistor. Technology is affecting our dailylife a lot.

document3 : You may find this tutorial on transistor technology interesting.

For explanation, let us treat the corpus as three documents and use simple token matching. In real Spark jobs, tokenization, case conversion, punctuation handling, stop-word removal, and stemming can change the final tokens and therefore the TF-IDF values.

TFIDF(technology, document2, corpus)

TF(technology, document2) = 2

IDF(technology, document2) = log((3+1)/(3+1)) = 0

TFIDF(technology, document2, corpus) = TF(technology, document2) . IDF(technology, document2) = 1 * 0 = 0

Even when the term ‘technology’ appeared twice in document2, as term has occurred in all the documents, it got no importance to document2 in the corpus.

TFIDF(TutorialKart, document1, corpus)

TF(TutorialKart, document1) = 1

IDF(TutorialKart, document1) = log((3+1)/(1+1)) = 1 (let’s take base as 2)

TFIDF(TutorialKart, document1, corpus) = TF(TutorialKart, document1) . IDF(TutorialKart, document1) = 1 * 1 = 1

The term ‘TutorialKart’ provides a TFIDF of 1.0 to document1 in the corpus.

How Spark MLlib performs TF-IDF with HashingTF and IDF

In Spark MLlib, TF and IDF are implemented separately.

Term frequency vectors could be generated using HashingTF or CountVectorizer.
IDF is an Estimator which is fit on a dataset and produces an IDFModel. The IDFModel takes feature vectors (generally created from HashingTF or CountVectorizer) and scales each column. Intuitively, it down-weights columns which appear frequently in a corpus.

The usual Spark ML TF-IDF pipeline is:

Create a DataFrame where each row represents one document.
Use Tokenizer or another tokenizer to convert each document into an array of words.
Use HashingTF or CountVectorizer to convert words into raw term-frequency vectors.
Fit IDF on the term-frequency vectors to learn document-frequency statistics from the corpus.
Use the resulting IDFModel to produce the final TF-IDF feature vectors.

HashingTF vs CountVectorizer for Spark TF-IDF

HashingTF uses the hashing trick. It maps tokens to a fixed number of feature indices. This is fast and does not require learning a vocabulary, but different terms can map to the same index when collisions occur. Increasing numFeatures reduces the chance of collisions.

CountVectorizer learns a vocabulary from the corpus and then creates term-frequency vectors using that vocabulary. It is useful when you want interpretable feature indices, vocabulary control, and options such as minimum document frequency. For small examples, either option works. For production text pipelines, choose based on vocabulary size, interpretability needs, and memory constraints.

Important preprocessing before Spark TF-IDF

TF-IDF quality depends heavily on the tokens provided to HashingTF or CountVectorizer. Before generating features, consider whether your use case needs lowercase conversion, punctuation cleanup, stop-word removal, n-grams, or custom tokenization. For example, Technology, technology, and technology. may become different tokens unless preprocessing handles case and punctuation consistently.

Spark MLlib TF-IDF – Java Example

In the following example, we will write a Java program to perform TF-IDF of documents. We are taking documents as list of strings.

TFIDFExample.java

</>

Copy

import java.util.Arrays;
import java.util.List;

import org.apache.spark.ml.feature.HashingTF;
import org.apache.spark.ml.feature.IDF;
import org.apache.spark.ml.feature.IDFModel;
import org.apache.spark.ml.feature.Tokenizer;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.RowFactory;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.Metadata;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;

public class TFIDFExample {
	public static void main(String[] args){

		// create a spark session
		SparkSession spark = SparkSession
				.builder()
				.appName("TFIDF Example")
				.master("local[2]")
				.enableHiveSupport()
				.getOrCreate();

		// documents corpus. each row is a document.
		List<Row> data = Arrays.asList(
				RowFactory.create(0.0, "Welcome to TutorialKart."),
				RowFactory.create(0.0, "Learn Spark at TutorialKart."),
				RowFactory.create(1.0, "Spark Mllib has TF-IDF.")
				);
		StructType schema = new StructType(new StructField[]{
				new StructField("label", DataTypes.DoubleType, false, Metadata.empty()),
				new StructField("sentence", DataTypes.StringType, false, Metadata.empty())
		});
		
		// import data with the schema
		Dataset<Row> sentenceData = spark.createDataFrame(data, schema);

		// break sentence to words
		Tokenizer tokenizer = new Tokenizer().setInputCol("sentence").setOutputCol("words");
		Dataset<Row> wordsData = tokenizer.transform(sentenceData);
		
		// define Transformer, HashingTF
		int numFeatures = 32;
		HashingTF hashingTF = new HashingTF()
				.setInputCol("words")
				.setOutputCol("rawFeatures")
				.setNumFeatures(numFeatures);

		// transform words to feature vector
		Dataset<Row> featurizedData = hashingTF.transform(wordsData);

		System.out.println("TF vectorized data\n----------------------------------------");
		for(Row row:featurizedData.collectAsList()){
			System.out.println(row.get(3));
		}

		System.out.println(featurizedData.toJSON());

		// IDF is an Estimator which is fit on a dataset and produces an IDFModel
		IDF idf = new IDF().setInputCol("rawFeatures").setOutputCol("features");
		IDFModel idfModel = idf.fit(featurizedData);

		// The IDFModel takes feature vectors (generally created from HashingTF or CountVectorizer) and scales each column
		Dataset<Row> rescaledData = idfModel.transform(featurizedData);
		
		System.out.println("TF-IDF vectorized data\n----------------------------------------");
		for(Row row:rescaledData.collectAsList()){
			System.out.println(row.get(4));
		}
		
		System.out.println("Transformations\n----------------------------------------");
		for(Row row:rescaledData.collectAsList()){
			System.out.println(row);
		}

		spark.close();
		
	}
}

Output

TF vectorized data
----------------------------------------
(32,[20,28,29],[1.0,1.0,1.0])
(32,[1,4,15,28],[1.0,1.0,1.0,1.0])
(32,[1,2,4,9],[1.0,1.0,1.0,1.0])


TF-IDF vectorized data
----------------------------------------
(32,[20,28,29],[0.6931471805599453,0.28768207245178085,0.6931471805599453])
(32,[1,4,15,28],[0.28768207245178085,0.28768207245178085,0.6931471805599453,0.28768207245178085])
(32,[1,2,4,9],[0.28768207245178085,0.6931471805599453,0.28768207245178085,0.6931471805599453])

Transformations
----------------------------------------
[0.0,Welcome to TutorialKart.,WrappedArray(welcome, to, tutorialkart.),(32,[20,28,29],[1.0,1.0,1.0]),(32,[20,28,29],[0.6931471805599453,0.28768207245178085,0.6931471805599453])]
[0.0,Learn Spark at TutorialKart.,WrappedArray(learn, spark, at, tutorialkart.),(32,[1,4,15,28],[1.0,1.0,1.0,1.0]),(32,[1,4,15,28],[0.28768207245178085,0.28768207245178085,0.6931471805599453,0.28768207245178085])]
[1.0,Spark Mllib has TF-IDF.,WrappedArray(spark, mllib, has, tf-idf.),(32,[1,2,4,9],[1.0,1.0,1.0,1.0]),(32,[1,2,4,9],[0.28768207245178085,0.6931471805599453,0.28768207245178085,0.6931471805599453])]

The Java example uses Spark ML DataFrame-based transformers. Tokenizer produces the words column, HashingTF produces the rawFeatures column, and IDFModel produces the final features column. These final vectors are the TF-IDF representation of the input documents.

PySpark TF-IDF example using HashingTF and IDF

Now, let us implement the same use case in Python.

spark-mllib-tfidf.py

</>

Copy

from __future__ import print_function

from pyspark.ml.feature import HashingTF, IDF, Tokenizer
from pyspark.sql import SparkSession

if __name__ == "__main__":
    spark = SparkSession\
        .builder\
        .appName("TfIdf Example")\
        .getOrCreate()

    sentenceData = spark.createDataFrame([
        (0.0, "Welcome to TutorialKart."),
        (0.0, "Learn Spark at TutorialKart."),
        (1.0, "Spark Mllib has TF-IDF.")
    ], ["label", "sentence"])

    tokenizer = Tokenizer(inputCol="sentence", outputCol="words")
    wordsData = tokenizer.transform(sentenceData)

    hashingTF = HashingTF(inputCol="words", outputCol="rawFeatures", numFeatures=20)
    featurizedData = hashingTF.transform(wordsData)
    # alternatively, CountVectorizer can also be used to get term frequency vectors

    idf = IDF(inputCol="rawFeatures", outputCol="features")
    idfModel = idf.fit(featurizedData)
    rescaledData = idfModel.transform(featurizedData)

    rescaledData.select("label", "features").show()

spark.stop()

Execute the following command to submit this Python Application to Spark and run it.

$ spark-submit spark-mllib-tfidf.py

A typical output contains one row per document and a sparse vector in the features column. The vector format shows the vector size, non-zero indices, and corresponding TF-IDF weights.

+-----+--------------------+
|label|            features|
+-----+--------------------+
|  0.0|(20,[...],[...])     |
|  0.0|(20,[...],[...])     |
|  1.0|(20,[...],[...])     |
+-----+--------------------+

Using CountVectorizer instead of HashingTF for Spark TF-IDF

The previous examples use HashingTF. If you want Spark to build a vocabulary from the corpus before calculating TF-IDF, use CountVectorizer. This is helpful when you need to inspect which word belongs to which feature index.

</>

Copy

from pyspark.ml.feature import CountVectorizer, IDF, Tokenizer
from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("CountVectorizer TF-IDF").getOrCreate()

df = spark.createDataFrame([
    (0, "spark mllib tf idf example"),
    (1, "spark machine learning example"),
    (2, "tf idf feature extraction in spark")
], ["id", "sentence"])

tokenizer = Tokenizer(inputCol="sentence", outputCol="words")
words_df = tokenizer.transform(df)

cv = CountVectorizer(inputCol="words", outputCol="rawFeatures", vocabSize=1000, minDF=1.0)
cv_model = cv.fit(words_df)
featurized_df = cv_model.transform(words_df)

idf = IDF(inputCol="rawFeatures", outputCol="features")
idf_model = idf.fit(featurized_df)
result = idf_model.transform(featurized_df)

result.select("id", "words", "features").show(truncate=False)

spark.stop()

Use HashingTF when you prefer a fixed-size vector without fitting a vocabulary. Use CountVectorizer when the vocabulary itself is useful for debugging, interpretation, or feature selection.

Common Spark TF-IDF mistakes and fixes

Mistake	Why it matters	Fix
Using raw sentences directly with IDF	`IDF` expects vector input, not a string column.	Tokenize text first, then create `rawFeatures` using `HashingTF` or `CountVectorizer`.
Choosing a very small `numFeatures` in `HashingTF`	More terms may collide into the same feature index.	Use a larger feature size for real corpora and validate model quality.
Skipping text normalization	Case and punctuation can create separate tokens for the same word.	Apply consistent preprocessing before TF-IDF.
Fitting IDF on test data	This leaks corpus statistics from evaluation data.	Fit `IDF` on training data only, then transform validation or test data.
Expecting TF-IDF vectors to show original words with HashingTF	HashingTF uses hashed indices, not a learned vocabulary.	Use `CountVectorizer` when feature-to-word interpretation is required.

Where Spark MLlib TF-IDF fits in a machine learning pipeline

TF-IDF usually appears before a machine learning algorithm. After creating the features column, you can pass it to algorithms such as Logistic Regression, Naive Bayes, Linear SVM, KMeans, or other Spark ML estimators that accept feature vectors. For supervised learning, the DataFrame commonly contains at least a label column and a features column.

For reproducible pipelines, combine the stages in a Spark ML Pipeline. This keeps tokenization, feature extraction, IDF fitting, and model training together, reducing the chance of applying different transformations during training and prediction.

Spark MLlib TF-IDF FAQs

What is the difference between TF and IDF in Spark MLlib?

TF measures how often a term appears in a document. IDF measures how rare or common that term is across the full corpus. Spark first creates TF vectors with HashingTF or CountVectorizer, then applies IDF to down-weight terms that appear in many documents.

Should I use HashingTF or CountVectorizer for TF-IDF in Spark?

Use HashingTF when you want a fast, fixed-size representation and do not need to inspect the learned vocabulary. Use CountVectorizer when you want a corpus vocabulary and more interpretable feature indices.

Why are Spark TF-IDF vectors shown as sparse vectors?

Text data usually has many possible terms, but each document contains only a small subset of them. Spark stores TF-IDF features as sparse vectors to save memory and processing time.

Can TF-IDF be used for Spark text classification?

Yes. TF-IDF vectors can be used as input features for text classification algorithms in Spark ML. A common workflow is tokenizer, TF transformer, IDF transformer, and then a classifier such as Logistic Regression or Naive Bayes.

Why do frequent words get low TF-IDF scores?

Frequent words that appear in many documents have low IDF values. Because TF-IDF multiplies term frequency by inverse document frequency, common terms receive lower weights even if they occur several times in a document.

Editorial QA checklist for Spark MLlib TF-IDF examples

Confirm that every TF-IDF example has a clear document column, token column, raw feature vector column, and final feature vector column.
Check that new command-line examples use language-bash and output-only examples use the output class.
Verify that HashingTF examples explain possible hash collisions and the role of numFeatures.
When CountVectorizer is used, explain that it fits a vocabulary before producing term-frequency vectors.
Ensure the tutorial distinguishes Spark ML DataFrame-based APIs from older RDD-based MLlib examples.

Conclusion

In this Spark Tutorial, we have learnt about Spark Mllib TF-IDF, how to calculate TF and IDF, and how to realize TF-IDF using Spark MLlib HashingTF Transformer and IDF Estimator. We also looked at when to use CountVectorizer, how preprocessing affects the final vectors, and what mistakes to avoid while building Spark TF-IDF pipelines.

TutorialKart.com

Spark MLlib TF-IDF – Example

Spark MLlib TF-IDF

What is TF-IDF in Spark MLlib?

Manual TF-IDF calculation for a small Spark-style corpus

TFIDF(technology, document2, corpus)

TFIDF(TutorialKart, document1, corpus)

How Spark MLlib performs TF-IDF with HashingTF and IDF

HashingTF vs CountVectorizer for Spark TF-IDF

Important preprocessing before Spark TF-IDF

Spark MLlib TF-IDF – Java Example

PySpark TF-IDF example using HashingTF and IDF

Using CountVectorizer instead of HashingTF for Spark TF-IDF

Common Spark TF-IDF mistakes and fixes

Where Spark MLlib TF-IDF fits in a machine learning pipeline

Spark MLlib TF-IDF FAQs

What is the difference between TF and IDF in Spark MLlib?

Should I use HashingTF or CountVectorizer for TF-IDF in Spark?

Why are Spark TF-IDF vectors shown as sparse vectors?

Can TF-IDF be used for Spark text classification?

Why do frequent words get low TF-IDF scores?

Editorial QA checklist for Spark MLlib TF-IDF examples

Conclusion

Popular Courses

SAP

CRM

SAP Resources

Apache

GUI

Programming

Databases

Mobile

Linux

Web & Server

Testing

Learning