Sklearn fit vs predict, order of columns matters?

Sklearn

machine learning

data preprocessing

feature order

model training

Sklearn fit vs predict, order of columns matters?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Start Practicing Learn More

Introduction

Yes, column order matters in scikit-learn. The fit method learns parameters based on the column positions of the training data, not column names. If you pass columns in a different order during predict, the model applies the wrong learned weights to the wrong features, producing incorrect predictions without raising an error. This is because NumPy arrays (which sklearn uses internally) are position-indexed, not name-indexed. The fix is to ensure consistent column ordering between training and prediction, typically by using a Pipeline or saving the column order alongside the model.

The Problem

python

1import pandas as pd
2from sklearn.linear_model import LinearRegression
3
4# Training data
5train = pd.DataFrame({
6    'height': [170, 180, 160, 175],
7    'weight': [70, 85, 55, 75],
8    'age': [25, 30, 22, 28],
9})
10y_train = [65000, 75000, 50000, 70000]
11
12model = LinearRegression()
13model.fit(train[['height', 'weight', 'age']], y_train)
14
15# Prediction with WRONG column order
16test = pd.DataFrame({
17    'age': [27],
18    'height': [172],
19    'weight': [68],
20})
21
22# BUG: columns are [age, height, weight] but model expects [height, weight, age]
23prediction = model.predict(test[['age', 'height', 'weight']])
24# No error raised, but prediction is WRONG

Sklearn converts DataFrames to NumPy arrays internally, losing column names. The model uses column positions, so swapped columns produce silently wrong results.

Fix 1: Explicit Column Ordering

python

1# Define column order once
2FEATURE_COLUMNS = ['height', 'weight', 'age']
3
4# Training
5model.fit(train[FEATURE_COLUMNS], y_train)
6
7# Prediction — always use the same column list
8prediction = model.predict(test[FEATURE_COLUMNS])

Store the column order as a constant and reference it everywhere. This is the simplest and most reliable approach.

Fix 2: Use a Pipeline with ColumnTransformer

python

1from sklearn.pipeline import Pipeline
2from sklearn.compose import ColumnTransformer
3from sklearn.preprocessing import StandardScaler
4from sklearn.linear_model import LinearRegression
5
6preprocessor = ColumnTransformer(
7    transformers=[
8        ('num', StandardScaler(), ['height', 'weight', 'age']),
9    ]
10)
11
12pipeline = Pipeline([
13    ('preprocessor', preprocessor),
14    ('model', LinearRegression()),
15])
16
17# Fit — column names are locked by ColumnTransformer
18pipeline.fit(train, y_train)
19
20# Predict — DataFrame with any column order works correctly
21test_shuffled = test[['weight', 'age', 'height']]  # Different order
22prediction = pipeline.predict(test_shuffled)  # Correct result

ColumnTransformer selects columns by name, not position, so the order of columns in the input DataFrame does not matter.

Fix 3: Save Column Order with the Model

python

1import joblib
2
3# Save model and column order together
4FEATURE_COLUMNS = ['height', 'weight', 'age']
5model.fit(train[FEATURE_COLUMNS], y_train)
6
7joblib.dump({
8    'model': model,
9    'columns': FEATURE_COLUMNS,
10}, 'model_bundle.pkl')
11
12# Load and use
13bundle = joblib.load('model_bundle.pkl')
14loaded_model = bundle['model']
15columns = bundle['columns']
16
17prediction = loaded_model.predict(test[columns])

When deploying a model, always serialize the expected column order alongside the model. This prevents column mismatch in production.

Fix 4: set_output API (Sklearn 1.2+)

python

1from sklearn.linear_model import LinearRegression
2
3model = LinearRegression()
4model.set_output(transform="pandas")
5
6# Sklearn 1.2+ raises a warning if column names differ
7model.fit(train[['height', 'weight', 'age']], y_train)
8
9# With feature_names_in_ attribute
10print(model.feature_names_in_)
11# array(['height', 'weight', 'age'], dtype=object)
12
13# This will raise a warning about mismatched feature names
14prediction = model.predict(test[['age', 'height', 'weight']])
15# UserWarning: X has feature names ['age', 'height', 'weight'],
16# but LinearRegression was fitted with ['height', 'weight', 'age']

Sklearn 1.2+ stores feature_names_in_ during fit and raises a warning (not an error) when predict receives different feature names or ordering.

Verifying Column Consistency

python

1# Check expected features after fitting
2print(model.feature_names_in_)  # ['height', 'weight', 'age']
3
4# Reorder test data to match training order
5test_reordered = test[model.feature_names_in_]
6prediction = model.predict(test_reordered)
7
8# Defensive function
9def safe_predict(model, X):
10    if hasattr(model, 'feature_names_in_'):
11        X = X[model.feature_names_in_]
12    return model.predict(X)

Common Pitfalls

Assuming column names are used internally: Sklearn converts DataFrames to NumPy arrays, discarding column names. Column position is all that matters to the model.
Different column order in train vs test splits: If you select features differently (e.g., train has df[['a','b']] but test has df[['b','a']]), predictions are silently wrong with no error.
Adding or removing features between fit and predict: If you train with 5 features but predict with 4, sklearn raises a ValueError. But if you predict with 5 different features in the same order, it runs without error.
Ignoring the feature_names_in_ warning: Sklearn 1.2+ warns about mismatched names but does not raise an error. Treat these warnings as errors in production by using warnings.filterwarnings('error').
Not saving column order when serializing models: Saving only the model object with joblib.dump(model) loses the expected column order. Always save the feature list or use a Pipeline.

Summary

Column order matters in sklearn — fit and predict must receive features in the same positional order
Define feature columns as a constant and use it for both training and prediction
Use ColumnTransformer in a Pipeline for name-based column selection that ignores order
Save the column list alongside the model when serializing for deployment
Check model.feature_names_in_ (sklearn 1.2+) to verify expected feature names and order