Avoiding Data Leakage with scikit‑learn Pipelines and ColumnTransformers

Problem

When building a predictive model on tabular data that mixes numeric and categorical columns, a common pitfall is to perform preprocessing before splitting the data into training and test sets. If a transformer learns statistics (e.g., mean, variance, category levels) from the full dataset, those statistics leak into the test set, inflating performance and giving a false sense of model quality.

Why Pipelines Matter

A Pipeline chains a sequence of transformers and an estimator into a single object. The fit method of the pipeline runs each transformer on the training data only, then fits the estimator on the transformed data. When you call predict on new data, the same transformers are applied in the exact same order, ensuring consistency between training and inference.

The ColumnTransformer extends this idea to mixed‑type data. It lets you specify separate preprocessing pipelines for numeric columns (e.g., scaling) and categorical columns (e.g., one‑hot encoding). The output is a unified feature matrix that can be fed directly into the estimator.

Building a Mixed‑Type Pipeline

# Example: Logistic regression on a toy dataset
import pandas as pd
from sklearn.model_selection import train_test_split, cross_validate
from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.linear_model import LogisticRegression

# Synthetic data
X = pd.DataFrame({
    'age': [25, 45, 35, 23, 52],
    'salary': [50000, 80000, 60000, 42000, 95000],
    'city': ['NY', 'SF', 'NY', 'LA', 'SF'],
    'gender': ['M', 'F', 'F', 'M', 'F']
})
y = pd.Series([0, 1, 0, 0, 1])

# Train / test split
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

# Define column-specific transformers
numeric_tf = Pipeline([('scaler', StandardScaler())])
categorical_tf = Pipeline([('onehot', OneHotEncoder(handle_unknown='ignore'))])

# Combine into a ColumnTransformer
preprocessor = ColumnTransformer(
    transformers=[
        ('num', numeric_tf, ['age', 'salary']),
        ('cat', categorical_tf, ['city', 'gender'])
    ],
    remainder='drop',  # drop any remaining columns
    sparse_threshold=0.0  # force dense output if estimator needs it
)

# Full pipeline: preprocessing + model
clf = Pipeline([
    ('prep', preprocessor),
    ('model', LogisticRegression(max_iter=200))
])

# Fit and evaluate
clf.fit(X_train, y_train)
print('Test accuracy:', clf.score(X_test, y_test))

# Inspect feature names
print('Feature names after ColumnTransformer:', preprocessor.get_feature_names_out())
print('Feature names after full Pipeline:', clf.get_feature_names_out())

Key points in the code:

• The StandardScaler learns the mean and std only from X_train.
• The OneHotEncoder learns the category mapping from X_train and ignores unseen categories in X_test.
• Setting sparse_threshold=0.0 forces a dense output, which is required for some estimators that do not accept sparse matrices.
• By calling clf.score we automatically apply the same preprocessing to X_test.

Checking for Leakage

To verify that the pipeline prevents leakage, compare two scenarios:

1. Global preprocessing: fit StandardScaler and OneHotEncoder on the entire dataset before splitting.
2. Pipeline preprocessing: fit the pipeline on the training split only.

Run cross_validate on both pipelines and observe the test scores. In the global case you’ll typically see inflated scores because the test set statistics influence the training transforms. The pipeline case should yield more realistic performance.

Trade‑offs and Limitations

• Sparse output: OneHotEncoder returns a sparse matrix by default. Some estimators (e.g., LinearRegression) expect dense arrays, so you must set sparse_threshold=0.0 or convert the output manually.
• Feature name tracking: get_feature_names_out only works reliably if all steps support the set_output API and you’re on scikit‑learn ≥0.24. Earlier versions may drop names.
• Custom transformers: The pipeline will not guard against transformers that inadvertently use target information (e.g., a transformer that encodes the target as a feature). You must design such steps carefully.
• Version differences: The behavior of ColumnTransformer regarding sparse output and column ordering can vary between releases. Always check the documentation for the exact version you’re using.

Takeaway

Using Pipeline with ColumnTransformer is the recommended pattern for mixed‑type tabular data in scikit‑learn. It guarantees that preprocessing is learned only on training data, eliminates leakage, and keeps the code tidy. Remember to:

1. Set sparse_threshold=0.0 if your estimator can’t handle sparse matrices.
2. Use set_output('pandas') when you want to preserve column names throughout the pipeline.
3. Validate the pipeline inside cross‑validation or grid search to ensure transforms are applied per fold.
4. Inspect get_feature_names_out to confirm that the feature mapping matches your expectations.

With these practices, you’ll build reproducible, leakage‑free models that scale from prototypes to production.