A Two-Stage Workflow to Detect Data Leakage.

Data leakage is the silent killer of machine learning models. You finish training, see an incredible 98% accuracy on your test set, and feel like a genius, until the model hits production and completely fails.

The most overlooked part of machine learning is the active prevention of data leakage. While robust cross-validation is a great way to verify if your model generalizes, it merely tells you that a problem exists. It doesn’t tell you where the leak is coming from.

To solve this, I’ve created a standardized, two-stage workflow designed to identify the exact columns causing data leakage directly and immediately. Before you spend hours tuning hyperparameters, run your data through these two lightning-fast checks.

Our first line of defense looks for direct, linear relationships. If a single column can perfectly (or almost perfectly) predict your target variable on its own, you likely have a leak—like an accidental "future" timestamp or an ID column that encodes the target label.

print("=== 1. Linear Leakage Check (Correlation & R²) ===")

correlations = x_train.corrwith(pd.Series(np.ravel(y_train), index=x_train.index), numeric_only=True).abs()
r2_approximations = correlations ** 2

leakage_df = pd.DataFrame({
    'Correlation (|r|)': correlations,
    'Single-Feature R²': r2_approximations
}).sort_values('Single-Feature R²', ascending=False)

high_leakage = leakage_df[leakage_df['Single-Feature R²'] > 0.10]

if len(high_leakage) > 0:
    print(f"🔴 Found {len(high_leakage)} potential leakage sources (R² > 10%):")
    print(high_leakage.head(20))
else:
    print("✅ No linearly leaky features found.")

(.abs())

because a strong negative correlation is just as predictive as a strong positive one.This step is computationally cheap. By relying on pandas' corrwith

method, the operations are heavily vectorized in C, meaning it can scan thousands of columns in a fraction of a second. It immediately highlights glaring linear leaks without requiring you to train a single model.

Linear correlation is great, but it misses complex, non-linear relationships. For example, a categorical ID mapped to integers might have a low linear correlation but still perfectly separate the target classes. To catch these, we move to Stage 2.

print("\n=== 2. Non-Linear Leakage Check (Tree Impurity) ===")

from sklearn.tree import DecisionTreeRegressor
import pandas as pd

dt = DecisionTreeRegressor(max_depth=5, random_state=42)
dt.fit(x_train, y_train)

tree_importance = pd.DataFrame({
    'Feature': x_train.columns,
    'Importance': dt.feature_importances_
}).sort_values('Importance', ascending=False)

top_tree_features = tree_importance[tree_importance['Importance'] > 0.10]

if len(top_tree_features) > 0:
    print(f"🔴 Found {len(top_tree_features)} suspicious features dominating tree splits:")
    print(top_tree_features.to_string(index=False))
else:
    print("✅ No single feature dominates the tree decisions.")

Train a Shallow Tree: We initialize a single DecisionTreeRegressor

with a shallow depth (max_depth=5)

. A decision tree works by continuously splitting the data into groups based on the feature that best separates the target variable.

Extract Feature Importance: Once the tree is fitted, we pull the feature_importances_

attribute. Scikit-learn calculates this based on impurity reduction—essentially tracking how much a specific feature helped the tree make clean splits.

Filter and Flag: We map these importances

back to the column names, sort them, and flag any feature that dictates more than 10% of the tree's overall decision-making process.

A single decision tree restricted to a depth of 5 is incredibly fast to train. If a specific feature is a massive source of non-linear data leakage, the tree is going to be instantly drawn to it, splitting on that feature at the very top (root) nodes. Furthermore, extracting the built-in impurity importance requires zero extra computation time once the model is fitted.

=== 1. Linear Leakage Check (Correlation & R²) ===
🔴 Found 6 potential leakage sources (R² > 10%):
                                                 Correlation (|r|)  \
target_encoder__apr_drg_code                              0.586974   
remainder__total_charges                                  0.578617   
target_encoder__ccs_diagnosis_code                        0.488559   
target_encoder__ccs_procedure_code                        0.480237   
target_encoder__apr_mdc_code                              0.376924   
one_hot_encoder__apr_severity_of_illness_code_4           0.347295   

                                                 Single-Feature R²  
target_encoder__apr_drg_code                              0.344538  
remainder__total_charges                                  0.334797  
target_encoder__ccs_diagnosis_code                        0.238690  
target_encoder__ccs_procedure_code                        0.230628  
target_encoder__apr_mdc_code                              0.142072  
one_hot_encoder__apr_severity_of_illness_code_4           0.120614  

=== 2. Non-Linear Leakage Check (Tree Impurity) ===

🔴 Found 2 suspicious features dominating tree splits:
                     Feature  Importance
    remainder__total_charges    0.752467
target_encoder__apr_drg_code    0.109771

By combining these two steps, you create an end-to-end, standardized workflow that catches both linear and non-linear data leakage before you commit to heavy model training. Embed this snippet at the top of your ML pipelines, and you'll never be caught off-guard by a leaky dataset in production again.