Handling Missing Data in ML Pipelines

9 min read

Learn

Quick Reading

Estimated 9 mins

Prereq

Intermediate

Basic ML concepts helpful

Interactive

Static Playbook

Static guide & reference tables

Missingness isn't random by default — check the mechanism first

Before choosing an imputation strategy, it matters *why* data is missing. Statisticians distinguish three mechanisms (Rubin, 1976): MCAR (missing completely at random — missingness is unrelated to any variable, observed or not), MAR (missing at random — missingness depends on *other observed* variables, e.g. income is more often missing for older respondents, but not on income itself), and MNAR (missing not at random — missingness depends on the unobserved value itself, e.g. high earners are specifically less likely to report income). Most imputation methods implicitly assume MCAR or MAR; under MNAR, even sophisticated imputation can introduce systematic bias, because the very fact that a value is missing carries information the imputation method can't see.

Simple imputation strategies

Mean, median, or mode imputation fills missing values with a single summary statistic of the observed values in that column. It's fast and simple but artificially reduces variance (every imputed row gets the identical value) and can distort relationships between features, especially when missingness is substantial or correlated with other variables. Median is generally preferred over mean for skewed or outlier-prone numeric features, for the same reason median is preferred elsewhere — it isn't dragged by extreme values.

imputation_strategies.py

python

import pandas as pd
import numpy as np
from sklearn.impute import SimpleImputer, KNNImputer
from sklearn.experimental import enable_iterative_imputer
from sklearn.impute import IterativeImputer

X = pd.DataFrame({
    "age": [25, np.nan, 40, 35, np.nan],
    "income": [50000, 62000, np.nan, 71000, 48000],
})

# 1. Simple: median imputation
simple = SimpleImputer(strategy="median")
X_simple = simple.fit_transform(X)

# 2. KNN imputation: fill using the average of the k most similar rows
knn = KNNImputer(n_neighbors=3)
X_knn = knn.fit_transform(X)

# 3. Iterative (MICE-style): model each missing column as a function
#    of the other columns, iterating until estimates stabilize
iterative = IterativeImputer(max_iter=10, random_state=0)
X_iter = iterative.fit_transform(X)

# 4. Always consider adding a missingness indicator alongside imputed values:
X["age_was_missing"] = X["age"].isna().astype(int)

KNN and iterative (MICE-style) imputation

KNN imputation fills a missing value using the average (or weighted average) of that feature's values from the $k$ most similar complete rows, which captures some of the correlation structure between features that simple imputation ignores — at the cost of being more computationally expensive and itself sensitive to feature scaling, since it relies on distance. Multiple Imputation by Chained Equations (MICE), implemented as IterativeImputer in scikit-learn, models each feature with missing values as a function of all other features, iterating round-robin across columns until the imputed values stabilize. It generally produces more statistically sound imputations than simple or KNN methods for MAR data, since it explicitly models the conditional relationships between features rather than relying only on marginal statistics or raw distance.

The missingness pattern itself can be predictive

Whether a value is missing is sometimes as predictive as the value itself would have been — e.g., a loan applicant who leaves income blank may carry different default risk than one who reports it, independent of the true income value. Adding a binary 'was this value missing' indicator column alongside the imputed value lets models use that signal explicitly, rather than having it disappear once imputation fills in a plausible-looking number.

Method	Captures feature correlations?	Computational cost	Best for
Mean/median/mode	No	Very low	Low missingness, quick baselines
KNN imputation	Partially (via distance)	Moderate-high	Moderate missingness, scaled features
Iterative/MICE	Yes (explicitly modeled)	High	MAR data, moderate-to-high missingness
Drop rows/columns	N/A	None	Very low missingness, or column is unusable

Fit imputers on training data only

The same leakage discipline that applies to feature scaling applies here: any imputer's statistics (mean, median, KNN neighbor computation, iterative model) must be fit exclusively on the training fold and applied — not refit — to validation and test data. Fitting on the combined dataset leaks test-set distributional information into training-time imputation.

What's Next?

Missing data is one source of unreliable features; outliers are a related but distinct data-quality issue, and the two interact — naive imputation can be badly thrown off by outliers in the observed portion of a column.

I build these systems professionally.

Whether it's a RAG pipeline, analytics migration, or AI workflow — let's talk.

Work With Me →