Using scikit‑learn Pipeline with ColumnTransformer to Avoid Data Leakage

Problem and Takeaway

When preparing data for a supervised model, any statistic (mean, variance, encoding mapping) that is derived from the whole dataset and then applied to the training set leaks information from the test fold into the model. This leads to overly optimistic performance estimates and models that fail in production. The takeaway is to encapsulate all preprocessing steps inside a sklearn.pipeline.Pipeline that is fitted only on the training data, using a sklearn.compose.ColumnTransformer to handle heterogeneous column types.

Requirements

• Input data may contain a mix of numeric and categorical columns.Each column type needs its own preprocessing (imputation, scaling, encoding).Preprocessing must be learned exclusively from the training split to avoid leakage.The workflow should remain reproducible across runs and environments.The final artifact should be a single estimators‑like object that can be persisted with joblib or pickle.

Smallest Suitable Design

The minimal architecture that satisfies the requirements is:

1. A ColumnTransformer that defines two branches:
• Numeric branch: SimpleImputer(strategy='median') → StandardScaler()
• Categorical branch: SimpleImputer(strategy='most_frequent') → OneHotEncoder(handle_unknown='ignore')
2. The ColumnTransformer is placed as the first step of a Pipeline.
3. The final step of the pipeline is the estimator (e.g., LogisticRegression(max_iter=1000)).

In code:

from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.linear_model import LogisticRegression

numeric_features = ['age', 'salary', 'loan_amount']
categorical_features = ['gender', 'education', 'property_area']

numeric_transformer = Pipeline(steps=[
    ('imputer', SimpleImputer(strategy='median')),
    ('scaler', StandardScaler())
])

categorical_transformer = Pipeline(steps=[
    ('imputer', SimpleImputer(strategy='most_frequent')),
    ('onehot', OneHotEncoder(handle_unknown='ignore'))
])

preprocess = ColumnTransformer(
    transformers=[
        ('num', numeric_transformer, numeric_features),
        ('cat', categorical_transformer, categorical_features)
    ])

model = Pipeline(steps=[
    ('preprocess', preprocess),
    ('classifier', LogisticRegression(max_iter=1000))
])

Trust and Data Boundaries

The pipeline enforces a clear boundary:

• When model.fit(X_train, y_train) is called, the ColumnTransformer learns imputation values, scaling parameters, and encoding mappings **only** from X_train.
• The same learned parameters are then applied to X_test via model.predict or model.transform without any re‑fitting.
• Because the preprocessing steps are hidden inside the pipeline, there is no external reference to the intermediate transformers that could be accidentally re‑fit on the full dataset.

Operational Checks

After fitting, you can verify that the pipeline behaved as expected:

• Inspect the fitted transformers: model.named_steps['preprocess'].transformers_ shows each branch and the underlying fitted steps.
• Check shapes: X_train_transformed = model.named_steps['preprocess'].transform(X_train) should have the same number of rows as X_train and a column count equal to the sum of numeric columns plus the one‑hot encoded categories.
• Run a quick cross‑validation loop (e.g., cross_val_score(model, X, y, cv=5)) and ensure scores are stable across folds; large variance may indicate leakage or data issues.
• If using scikit‑learn ≥ 0.24, confirm sparse output: issparse(X_train_transformed) should return True when the categorical branch uses the default sparse OneHotEncoder.

Failure Modes and Conditions That Would Change the Design

• New column types: If text or image features appear, the current two‑branch ColumnTransformer is insufficient. You would add a third branch (e.g., TfidfVectorizer) or replace the categorical encoder with a hashing trick to keep dimensionality bounded.
• Very high‑cardinality categorical variables: A dense OneHotEncoder can explode memory. Switch to OneHotEncoder(sparse=True) (default) or use FeatureHasher to impose a fixed output size.
• Large‑scale data that does not fit in memory: The in‑memory fit of StandardScaler may fail. Replace it with StandardScaler(with_mean=False) that works on sparse matrices, or adopt an incremental learner such as SGDClassifier with partial_fit.
• Estimator requiring raw NumPy arrays: Some custom models expect a plain ndarray and reject sparse input. In that case, set sparse_threshold=0 in the ColumnTransformer to force dense output, or place a FunctionTransformer that converts sparse to dense after preprocessing.
• Changing scikit‑learn version: Versions prior to 0.24 did not preserve sparse output by default. If you rely on sparsity, either upgrade or explicitly set sparse_threshold to a low value.

Concrete Example: Verification of No Leakage

To demonstrate that the pipeline does not leak information, you can perform the following sanity check (run in a Python script or interactive session with read access to your dataset):

import numpy as np
from sklearn.model_selection import train_test_split

# Assume X, y are loaded
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Fit the pipeline
model.fit(X_train, y_train)

# Extract the learned imputation means for numeric columns
num_imp = model.named_steps['preprocess'].named_transformers_['num'].named_steps['imputer']
print('Learned imputation values:', num_imp.statistics_)

# Compute the mean of the same columns in the full dataset (should differ)
full_mean = np.nanmean(X[num_imp.feature_names_in_], axis=0)
print('Full‑dataset mean:', full_mean)

# Assert they are not identical (allowing for tiny floating‑point differences)
assert not np.allclose(num_imp.statistics_, full_mean, rtol=1e-10), \
    'Imputation statistics match the full dataset – possible leakage'

If the assertion passes, you have evidence that the fitting step used only the training partition. No output is invented here; you should see the printed arrays and the assertion succeed when the pipeline is correctly configured.

Practical Way to Check the Result

After training, persist the pipeline and later reload it to ensure reproducibility:

import joblib
joblib.dump(model, 'model_pipeline.joblib')
...
model_loaded = joblib.load('model_pipeline.joblib')
# Predict on unseen data
preds = model_loaded.predict(X_new)

Verify that preds has the same shape as the number of rows in X_new and contains only the classes expected by the estimator.

Conclusion
The described architecture — a Pipeline wrapping a ColumnTransformer — provides the smallest reliable design for preprocessing heterogeneous data while guarding against leakage. Operational checks such as inspecting fitted transformers, verifying shapes, and running cross‑validation give confidence that the pipeline behaves as intended. When data characteristics evolve (new column types, extreme cardinality, or volume limits), the design can be extended by adding branches, switching to sparse or hashing encoders, or adopting out‑of‑core learners, all without breaking the core principle that fitting occurs exclusively on training data.