DEV Community

Cover image for Leakage in ML Pipelines: How to build a bulletproof preprocessing architecture
Pasquale Molinaro
Pasquale Molinaro

Posted on • Originally published at Medium

Leakage in ML Pipelines: How to build a bulletproof preprocessing architecture

You build a model, run the evaluations, and hit a 95% accuracy on your test set. You deploy it to production feeling like a genius, only to watch it fail miserably on real-world data. We’ve all been there. When a model explodes in production after perfect local testing, the culprit is rarely the algorithm itself. Most of the time, it’s a silent architecture flaw introduced during the very first steps of preprocessing: Data Leakage. In this article, we will discover how one of the most common mistakes in handling Missing Values and oversampling implicitly corrupts your test data, and how to build a bulletproof, leak-free pipeline using Scikit-Learn.

The common error is preprocessing before splitting
Let’s look at a classic approach to data preparation. If you have a dataset with missing values and a mix of categorical and numerical columns, the most intuitive approach is to clean everything up before feeding it to the model.

You split the dataframes by type, apply a LabelEncoder to the text, and use an Imputer to fill the NaNs. The code usually looks something like this:

import pandas as pd
import numpy as np
from sklearn.preprocessing import LabelEncoder


# 1. Splitting categorical and numerical data
df_num = df_copy.select_dtypes(include=[np.number])
df_cat = df_copy.select_dtypes(include=['object'])

# 2. Encoding categorical data
for attr in df_cat.columns:
    df_cat[attr] = LabelEncoder().fit_transform(df_cat[attr])

# 3. Merging back and imputing missing values (omitted for brevity)
df_encoded = pd.concat([df_cat, df_num], axis=1)

# 4. Splitting into Train and Test ONLY at the end
X = df_encoded.drop(['TargetVariable'], axis=1)
y = df_encoded['TargetVariable']
y = df_encoded['TargetVariable']X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
Enter fullscreen mode Exit fullscreen mode

It looks logical, clean, and perfectly executed. But this code is a ticking time bomb. By applying .fit_transform() on the entire categorical dataframe before calling train_test_split, your encoder has just learned the global distribution of every single category, including the ones that will end up in your test set. If you apply an Imputer to fill missing values using the mean or median of the entire column, your training data is mathematically contaminated by the future testing data.

Your model isn’t predicting, it’s cheating.

The Leak-Proof Architecture

Step 1: Isolate the Test Set Immediately

The golden rule of production-ready machine learning is simple: quarantine your test data before you do anything else. Do not impute, do not encode, do not even look at it. You cut the dataset raw and dirty.

from sklearn.model_selection import train_test_split

# Define your target column name
TARGET_COL = 'target_variable' # e.g., 'RainTomorrow'

# Separate features and target from your raw dataframe
X = df.drop([TARGET_COL], axis=1)
y = df[TARGET_COL]

# 1. Split the raw, dirty data FIRST
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
Enter fullscreen mode Exit fullscreen mode

To put this in perspective, as an illustrative example: a model with data leakage might show a dazzling 98% accuracy during local testing. But when you fix the pipeline and evaluate it cleanly, the true baseline accuracy drops to 82%. That 16% gap is the illusion of data leakage.

Step 2: Define Independent Transformers
Now that our X_test is safely locked away, we need a way to clean X_train. But instead of manually hacking DataFrames and looping through columns, we define "Transformers". Think of them as blueprints for cleaning data. They don't do anything yet; they just hold the logic.

We will use Scikit-Learn’s ColumnTransformer to handle numerical and categorical columns independently, entirely avoiding messy pd.concat operations and index misalignments.

import numpy as np
from sklearn.compose import ColumnTransformer, make_column_selector as selector
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import OneHotEncoder
from sklearn.pipeline import Pipeline

# Blueprint for numbers: fill missing with median
numeric_transformer = SimpleImputer(strategy='median')

# Blueprint for categories: fill missing with most frequent, then encode
categorical_transformer = Pipeline(steps=[
    ('imputer', SimpleImputer(strategy='most_frequent')),
    ('encoder', OneHotEncoder(handle_unknown='ignore', sparse_output=False))
])

# Bundle the blueprints together using dynamic selection
preprocessor = ColumnTransformer(
    transformers=[
        ('num', numeric_transformer, selector(dtype_include=np.number)),
        ('cat', categorical_transformer, selector(dtype_include=object))
    ])
Enter fullscreen mode Exit fullscreen mode

Step 3: The Leak-Proof Pipeline in Action
Now we assemble the final architecture. Notice a crucial detail: we are importing the Pipeline from imblearn.pipeline, not the standard sklearn.pipeline. This is the secret sauce, and it’s where most junior developers get stuck with a TypeError. Why can't we just use the standard Scikit-Learn pipeline?

Under the hood, a standard Scikit-Learn pipeline expects every step (except the final estimator) to have a .transform() method. It assumes you are just modifying existing rows (like scaling numbers or encoding text). However, SMOTE doesn't just transform data; it generates entirely new synthetic rows using a method called .fit_resample(). If you put SMOTE inside a standard Scikit-Learn pipeline, it crashes. The imblearn pipeline is explicitly overridden to handle this structural difference. It knows exactly when to call .fit_resample() (only during training) and when to completely ignore the SMOTE step (during .predict() or .score()). By using this architecture, you guarantee that your test set remains pristine and untouched by synthetic data generation.

The Cross-Validation Trap
You might think you are safe if you use Cross-Validation (CV). But if you apply SMOTE or Standardization to your entire X_train before passing it to cross_val_score, you are leaking data across your folds! The beauty of the imblearn pipeline is that you can pass the entire pipeline object directly into Scikit-Learn’s cross_val_score. The framework will strictly apply SMOTE only to the training folds of each iteration, leaving the validation fold completely untouched.

from imblearn.pipeline import Pipeline as ImbPipeline
from imblearn.over_sampling import SMOTE
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import cross_val_score

# Assemble the ultimate leak-proof pipeline
pipeline = ImbPipeline(steps=[
    ('preprocessor', preprocessor),
    ('smote', SMOTE(random_state=42)),
    ('classifier', RandomForestClassifier(random_state=42)) 
])

# --- THE CROSS-VALIDATION PROOF ---
# The pipeline guarantees SMOTE is applied strictly to the training folds 
# of each split, never touching the validation folds.
cv_scores = cross_val_score(pipeline, X_train, y_train, cv=5)
print(f"Cross-Validation Accuracy: {cv_scores.mean():.4f} (+/- {cv_scores.std() * 2:.4f})")

# --- FINAL PRODUCTION EVALUATION ---
# Train everything in a single step on the Train Set ONLY
pipeline.fit(X_train, y_train)

# Evaluate on the Test Set safely. 
accuracy = pipeline.score(X_test, y_test)
print(f"True Production Accuracy: {accuracy:.4f}")
Enter fullscreen mode Exit fullscreen mode

The Leakage Iceberg
While imputation and categorical encoding are the most frequent culprits for junior developers, remember that Data Leakage has many faces. Calculating target-based features (like historical averages), applying a StandardScaler, or running feature selection algorithms (like PCA or statistical tests) on your entire dataset before splitting will also silently contaminate your test set. The Golden Rule applies to your entire pipeline: Split first, ask questions later.

Conclusion
Data Leakage is the silent killer of machine learning models. Splitting your raw dataset immediately and wrapping all your preprocessing logic into a strict pipeline isn’t just “clean code”, it is the only mathematical guarantee that your model’s performance in local testing will match its performance in production. Stop hacking DataFrames with manual .fit_transform() loops, and start engineering bulletproof pipelines. Your future self (and your production server) will thank you.

Here the full code:

import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.compose import ColumnTransformer, make_column_selector as selector
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import OneHotEncoder 
from sklearn.pipeline import Pipeline 
from imblearn.pipeline import Pipeline as ImbPipeline
from imblearn.over_sampling import SMOTE
from sklearn.ensemble import RandomForestClassifier

# --- 1. DATA LOADING & INITIAL SPLIT ---
df = pd.read_csv('your_dataset.csv')
TARGET_COL = 'target_column_name'

X = df.drop([TARGET_COL], axis=1)
y = df[TARGET_COL]

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# --- 2. PREPROCESSING LOGIC ---
numeric_transformer = SimpleImputer(strategy='median')

# We use standard sklearn Pipeline here, as no sampling happens at this stage
categorical_transformer = Pipeline(steps=[
    ('imputer', SimpleImputer(strategy='most_frequent')),
    ('encoder', OneHotEncoder(handle_unknown='ignore', sparse_output=False))
])

preprocessor = ColumnTransformer(
    transformers=[
        ('num', numeric_transformer, selector(dtype_include=np.number)),
        ('cat', categorical_transformer, selector(dtype_include=object))
    ])

# --- 3. FINAL PIPELINE ASSEMBLY ---
# Here we MUST use imblearn.pipeline to handle SMOTE safely
pipeline = ImbPipeline(steps=[
    ('preprocessor', preprocessor),
    ('smote', SMOTE(random_state=42)),
    ('classifier', RandomForestClassifier(random_state=42))
])

# --- 4. EXECUTION ---
pipeline.fit(X_train, y_train)
accuracy = pipeline.score(X_test, y_test)
print(f"Final Leak-Free Accuracy: {accuracy:.4f}")
Enter fullscreen mode Exit fullscreen mode

Top comments (0)