pandas
sklearn
CountVectorizer
data manipulation
Python programming

Insert result of sklearn CountVectorizer in a pandas dataframe

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

CountVectorizer returns a sparse matrix, not a pandas DataFrame. To put the result into a DataFrame, you need two extra pieces: the generated feature names and a conversion method that preserves the sparse structure when the dataset is large.

The modern and usually best approach is pd.DataFrame.sparse.from_spmatrix. It lets you inspect word counts in tabular form without forcing the entire matrix into a dense array.

Fit the Vectorizer and Get Feature Names

Start by fitting the vectorizer as usual.

python
1import pandas as pd
2from sklearn.feature_extraction.text import CountVectorizer
3
4texts = [
5    "blue car fast",
6    "red car slow",
7    "blue bike fast"
8]
9
10vectorizer = CountVectorizer()
11X = vectorizer.fit_transform(texts)
12feature_names = vectorizer.get_feature_names_out()
13
14print(X.shape)
15print(feature_names)

X is a sparse matrix where rows are documents and columns are vocabulary terms. feature_names provides the correct column labels in the same order as the matrix columns.

Convert the Sparse Matrix to a DataFrame

Use pandas' sparse constructor so you do not accidentally explode memory.

python
1df_counts = pd.DataFrame.sparse.from_spmatrix(
2    X,
3    columns=feature_names
4)
5
6print(df_counts)

That gives you a DataFrame where each column is a token and each row corresponds to one original text. For the sample data, you will see count columns like bike, blue, car, fast, red, and slow.

Preserve the Original Row Index

If your source texts already live in a DataFrame, preserve that index so later joins stay aligned.

python
1df = pd.DataFrame({
2    "doc_id": [101, 102, 103],
3    "text": texts,
4    "label": [1, 0, 1]
5})
6
7X = vectorizer.fit_transform(df["text"])
8df_counts = pd.DataFrame.sparse.from_spmatrix(
9    X,
10    index=df.index,
11    columns=vectorizer.get_feature_names_out()
12)
13
14print(df_counts)

This matters because feature matrices and metadata often get joined later. If row order changes and you did not carry the index through, labels can silently attach to the wrong documents.

Join Metadata and Count Features

Once the sparse DataFrame exists, combine it with the metadata columns you want to keep.

python
1result = pd.concat([
2    df[["doc_id", "label"]],
3    df_counts
4], axis=1)
5
6print(result)

This produces a single table that is easy to inspect during debugging, error analysis, or exploratory feature engineering.

When Dense Conversion Is Acceptable

For tiny datasets, you can convert to a dense NumPy array and then build a DataFrame from that.

python
1df_dense = pd.DataFrame(
2    X.toarray(),
3    columns=vectorizer.get_feature_names_out()
4)
5
6print(df_dense)

This is fine for toy examples and notebooks. It is a bad default for real corpora because dense arrays waste memory when most counts are zero.

Reuse the Same Vocabulary at Inference Time

If you are creating a DataFrame for model input, fit the vectorizer on training data and reuse it later with transform, not fit_transform again.

python
1train_texts = ["blue car fast", "red car slow"]
2test_texts = ["blue slow car"]
3
4vectorizer = CountVectorizer()
5X_train = vectorizer.fit_transform(train_texts)
6X_test = vectorizer.transform(test_texts)
7
8df_test = pd.DataFrame.sparse.from_spmatrix(
9    X_test,
10    columns=vectorizer.get_feature_names_out()
11)
12
13print(df_test)

If you refit on test data, the vocabulary can change and the columns will no longer line up with what your model expects.

Common Pitfalls

  • Using X.toarray() on a large sparse matrix and running into unnecessary memory pressure.
  • Forgetting get_feature_names_out(), which leaves you with unlabeled numeric columns.
  • Losing the original row index and then joining labels or metadata to the wrong rows.
  • Calling fit_transform on inference data, which changes the vocabulary and feature layout.
  • Assuming every token from raw text survives vectorization unchanged, even though stop words and token rules may remove some of them.

Summary

  • 'CountVectorizer returns a sparse matrix, so convert it with pandas' sparse DataFrame support when possible.'
  • Use get_feature_names_out() for correct column names.
  • Preserve row index if you plan to join features back to metadata.
  • Dense conversion is acceptable only for small examples and quick inspection.
  • Reuse the fitted vocabulary with transform to keep training and inference features aligned.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.