Practice Now
Sklearn
machine learning
data preprocessing
feature order
model training

Sklearn fit vs predict, order of columns matters?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design
Machine LearningData Science

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
Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design
Codemia
© 2026 Codemia
Resources
Legal
Social
All Rights Reserved.