scikit-learn
machine learning
SimpleImputer
CountVectorizer
data preprocessing

How to include SimpleImputer before CountVectorizer in a scikit-learn Pipeline?

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

Putting SimpleImputer before CountVectorizer is a common need when text data contains missing values. The catch is that SimpleImputer works with a two-dimensional input shape, while CountVectorizer expects a one-dimensional iterable of strings. The fix is to impute first and then flatten the column back into the shape that CountVectorizer expects.

Why the Shape Mismatch Happens

SimpleImputer is designed for tabular features, so it expects input shaped like rows and columns. CountVectorizer, by contrast, expects something like a list or series of text documents.

That means this idea is right in principle:

  • replace missing text with an empty string or placeholder
  • vectorize the cleaned text

But a direct pipeline fails unless you bridge the shape difference between the two steps.

Use FunctionTransformer to Flatten After Imputation

One clean solution is:

  1. select the text column
  2. impute missing values
  3. flatten from two-dimensional column shape to one-dimensional text sequence
  4. run CountVectorizer
python
1import pandas as pd
2from sklearn.feature_extraction.text import CountVectorizer
3from sklearn.impute import SimpleImputer
4from sklearn.pipeline import Pipeline
5from sklearn.preprocessing import FunctionTransformer
6
7df = pd.DataFrame(
8    {
9        "text": ["red apple", None, "green apple", "truck engine"],
10        "label": [0, 0, 0, 1],
11    }
12)
13
14X = df[["text"]]
15y = df["label"]
16
17text_pipeline = Pipeline(
18    [
19        ("imputer", SimpleImputer(strategy="constant", fill_value="")),
20        ("flatten", FunctionTransformer(lambda x: x.ravel(), validate=False)),
21        ("vectorizer", CountVectorizer()),
22    ]
23)
24
25X_transformed = text_pipeline.fit_transform(X, y)
26print(X_transformed.shape)

The ravel() call is the important bridge. It converts the imputed single-column array into the one-dimensional form required by the vectorizer.

Put It Inside a Full Modeling Pipeline

In real work, the text preprocessing usually feeds directly into a classifier.

python
1import pandas as pd
2from sklearn.feature_extraction.text import CountVectorizer
3from sklearn.impute import SimpleImputer
4from sklearn.linear_model import LogisticRegression
5from sklearn.pipeline import Pipeline
6from sklearn.preprocessing import FunctionTransformer
7
8df = pd.DataFrame(
9    {
10        "text": ["red apple", None, "green apple", "truck engine"],
11        "label": [0, 0, 0, 1],
12    }
13)
14
15X = df[["text"]]
16y = df["label"]
17
18model = Pipeline(
19    [
20        ("imputer", SimpleImputer(strategy="constant", fill_value="")),
21        ("flatten", FunctionTransformer(lambda x: x.ravel(), validate=False)),
22        ("vectorizer", CountVectorizer()),
23        ("classifier", LogisticRegression(max_iter=1000)),
24    ]
25)
26
27model.fit(X, y)
28print(model.predict(pd.DataFrame({"text": ["apple"]})))

This keeps missing-value handling, text vectorization, and modeling in one reproducible object.

When SimpleImputer Is Not Necessary

If you already have a pandas Series, a simpler option is often to fill missing values before the pipeline:

python
texts = df["text"].fillna("")

Then you can pass that series straight into CountVectorizer.

The reason to keep SimpleImputer in the pipeline is usually consistency:

  • cross-validation applies the same preprocessing automatically
  • production inference follows the same steps as training
  • preprocessing logic stays attached to the model

So both approaches are valid. The pipeline version is usually cleaner for deployable workflows.

Multiple Text Columns Need a Slightly Different Design

If you have several text columns, combine them first or use a ColumnTransformer with one text pipeline per column. The same shape rule still applies: each text branch must end up as a one-dimensional document sequence before CountVectorizer runs.

Trying to hand a full two-dimensional table directly to one CountVectorizer usually means the pipeline design still needs another transformation step.

Common Pitfalls

The first pitfall is forgetting that SimpleImputer outputs a two-dimensional array. CountVectorizer expects one-dimensional text input, so some form of flattening is required.

Another issue is using numeric imputation strategies on text. For missing text, strategy="constant" with an empty string is usually the sensible default.

Developers also move the missing-value cleanup outside the training pipeline and then forget to repeat it during inference. That creates training-serving mismatches.

Finally, if you already have a clean text Series, do not overcomplicate the design. The pipeline is useful when it solves a real consistency problem.

Summary

  • 'SimpleImputer can go before CountVectorizer, but you must fix the shape mismatch.'
  • The usual bridge is FunctionTransformer(lambda x: x.ravel(), validate=False).
  • Use a constant text fill value such as an empty string for missing documents.
  • Keeping imputation inside the pipeline improves training and inference consistency.
  • For multiple text columns, combine them first or build separate text branches with a ColumnTransformer.

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.