How to impute missing values with scikit-learn

Tabular model data often contains blanks from optional fields, delayed measurements, or partially collected records. In scikit-learn, those gaps need to be converted into feature values before most estimators can fit or predict without rejecting the input matrix.

SimpleImputer learns replacement values during fit() and applies them during transform(). Numeric columns can use a median or mean, while categorical columns usually use the most frequent value or a fixed category before encoding.

A mixed table should keep numeric and categorical imputers in separate ColumnTransformer branches. Placing that transformer inside a Pipeline keeps imputation fitted on training rows only, then reuses the same fill values when holdout rows or future records reach the estimator.

1. Create impute_missing_values.py with separate imputers for numeric and categorical columns.

   impute_missing_values.py

   import numpy as np
   import pandas as pd
   import sklearn
   from sklearn.compose import ColumnTransformer
   from sklearn.impute import SimpleImputer
   from sklearn.linear_model import LogisticRegression
   from sklearn.pipeline import Pipeline
   from sklearn.preprocessing import OneHotEncoder
    
    
   numeric_features = ["age", "monthly_spend"]
   categorical_features = ["plan"]
    
   train = pd.DataFrame(
       {
           "age": [25, 32, np.nan, 45, 51, 28, np.nan, 39],
           "monthly_spend": [120, 210, 180, np.nan, 320, 150, 170, 260],
           "plan": [
               "basic",
               "premium",
               "basic",
               np.nan,
               "premium",
               "basic",
               "standard",
               np.nan,
           ],
           "converted": [0, 1, 0, 1, 1, 0, 0, 1],
       }
   )
    
   holdout = pd.DataFrame(
       {
           "age": [np.nan, 42],
           "monthly_spend": [190, np.nan],
           "plan": [np.nan, "premium"],
       }
   )
    
   numeric_imputer = SimpleImputer(strategy="median")
   categorical_transformer = Pipeline(
       steps=[
           ("imputer", SimpleImputer(strategy="most_frequent")),
           ("encoder", OneHotEncoder(handle_unknown="ignore", sparse_output=False)),
       ]
   )
    
   preprocessor = ColumnTransformer(
       transformers=[
           ("numeric", numeric_imputer, numeric_features),
           ("categorical", categorical_transformer, categorical_features),
       ],
       verbose_feature_names_out=False,
   )
    
   model = Pipeline(
       steps=[
           ("preprocess", preprocessor),
           ("classifier", LogisticRegression(max_iter=1000, random_state=0)),
       ]
   )
    
   X_train = train[numeric_features + categorical_features]
   y_train = train["converted"]
    
   model.fit(X_train, y_train)
    
   fitted_preprocessor = model.named_steps["preprocess"]
   fitted_numeric = fitted_preprocessor.named_transformers_["numeric"]
   fitted_categorical = fitted_preprocessor.named_transformers_["categorical"]
   fitted_category_imputer = fitted_categorical.named_steps["imputer"]
    
   transformed_train = fitted_preprocessor.transform(X_train)
   transformed_holdout = fitted_preprocessor.transform(holdout)
   predictions = model.predict(holdout)
    
   print(f"scikit-learn {sklearn.__version__}")
   print("Input missing values:")
   print(X_train.isna().sum().to_string())
    
   print()
   print("Numeric fill values:")
   for name, value in zip(numeric_features, fitted_numeric.statistics_):
       print(f"- {name}: {value:.1f}")
    
   print()
   print("Categorical fill values:")
   for name, value in zip(categorical_features, fitted_category_imputer.statistics_):
       print(f"- {name}: {value}")
    
   print()
   print("Transformed feature names:")
   print(", ".join(fitted_preprocessor.get_feature_names_out()))
    
   print()
   print(
       "Transformed training shape: "
       f"{transformed_train.shape[0]} rows x {transformed_train.shape[1]} columns"
   )
   print(f"Missing values after transform: {int(np.isnan(transformed_train).sum())}")
    
   print()
   print("Transformed holdout rows:")
   holdout_table = pd.DataFrame(
       transformed_holdout,
       columns=fitted_preprocessor.get_feature_names_out(),
   ).round(2)
   print(holdout_table.to_string(index=False))
   print(f"Holdout predictions: {predictions.tolist()}")

   The numeric branch fills age and monthly_spend before model fitting. The categorical branch fills plan before OneHotEncoder expands it into binary features.

2. Run the imputation smoke script.

   $ python impute_missing_values.py
   scikit-learn 1.9.0
   Input missing values:
   age              2
   monthly_spend    1
   plan             2
   
   Numeric fill values:
   - age: 35.5
   - monthly_spend: 180.0
   
   Categorical fill values:
   - plan: basic
   
   Transformed feature names:
   age, monthly_spend, plan_basic, plan_premium, plan_standard
   
   Transformed training shape: 8 rows x 5 columns
   Missing values after transform: 0
   
   Transformed holdout rows:
    age  monthly_spend  plan_basic  plan_premium  plan_standard
   35.5          190.0         1.0           0.0            0.0
   42.0          180.0         0.0           1.0            0.0
   Holdout predictions: [1, 1]

3. Confirm the training input contains missing values before imputation.

   Input missing values:
   age              2
   monthly_spend    1
   plan             2

4. Check the fitted imputer statistics.

   Numeric fill values:
   - age: 35.5
   - monthly_spend: 180.0
   
   Categorical fill values:
   - plan: basic

   The medians and most frequent category are learned during fit(). Fit the pipeline on training rows only so validation or production rows do not change these values.

5. Verify the transformed matrix contains no missing values.

   Transformed feature names:
   age, monthly_spend, plan_basic, plan_premium, plan_standard
   
   Transformed training shape: 8 rows x 5 columns
   Missing values after transform: 0

   If missingness is itself a model signal, set add_indicator=True on the relevant imputer or add a MissingIndicator branch. Indicators are created for features that contain missing values during fit().

6. Confirm prediction accepts holdout rows that still contain missing input values.

   Transformed holdout rows:
    age  monthly_spend  plan_basic  plan_premium  plan_standard
   35.5          190.0         1.0           0.0            0.0
   42.0          180.0         0.0           1.0            0.0
   Holdout predictions: [1, 1]

7. Remove the scratch script after the smoke test.

   $ rm impute_missing_values.py