2.6 Comparing and Applying Mean, Median, and Mode Imputation

2.6 Comparing and Applying Mean, Median, and Mode Imputation

Missing data is one of the most pervasive data quality challenges in production machine learning. Null values can originate from faulty sensors, database join mismatches, user opt-outs on web forms, system timeouts, or upstream pipeline failures. Many popular machine learning algorithms (including standard scikit-learn estimators, linear models, and neural networks) will throw immediate runtime errors when presented with missing entries NaN or None.

To build robust, production-grade ML workflows on Databricks, machine learning practitioners must understand the theoretical mechanisms of missingness, the mathematical tradeoffs of various imputation strategies, and how to execute distributed imputation in PySpark without introducing data leakage.

---

Theoretical Taxonomy of Missing Data Mechanisms

In statistical theory (established by Donald Rubin), missing data is classified into three distinct mechanisms. Identifying the mechanism dictates whether dropping data introduces bias:

                          MISSING DATA MECHANISMS
                                     │
        ┌────────────────────────────┼────────────────────────────┐
        ▼                            ▼                            ▼
      MCAR                          MAR                          MNAR
Missing Completely at Random   Missing at Random         Missing Not at Random
• Independent of all data    • Depends on observed data   • Depends on the missing
• No bias if dropped          • Biased if dropped         • Highest risk of bias
• Purely random glitch        • Can impute via features   • Must use indicators

Missing Completely at Random (MCAR)

• Definition: The probability of a value being missing is completely independent of both observed variables and unobserved values ($P(M \mid Y_{obs}, Y_{mis}) = P(M)$).
• Real-World Example: A lab technician accidentally drops a random test tube, or a temporary network glitch drops arbitrary IoT sensor packets.
• Consequence: Deleting missing rows (Complete Case Analysis) reduces statistical power and sample size, but does not introduce systematic bias into model parameters.

Missing at Random (MAR)

• Definition: The probability of missingness depends systematically on other observed features in the dataset, but not on the unobserved missing value itself ($P(M \mid Y_{obs}, Y_{mis}) = P(M \mid Y_{obs})$).
• Real-World Example: Male survey respondents are statistically less likely to answer a question regarding depressive symptoms, but gender is recorded and observed.
• Consequence: Dropping rows introduces severe selection bias. However, because the missingness is explained by observed features, multivariate imputation or conditional statistical imputation can restore unbiased estimates.

Missing Not at Random (MNAR)

• Definition: The probability of missingness depends directly on the unobserved missing value itself ($P(M \mid Y_{obs}, Y_{mis}) \ne P(M \mid Y_{obs})$).
• Real-World Example: Individuals with extremely high annual incomes or heavy debt burdens decline to disclose their financial numbers on credit applications.
• Consequence: The most dangerous mechanism. Imputing standard means or medians suppresses variance and introduces extreme bias. Data scientists must engineer missingness indicator features to explicitly retain the signal that the value was withheld.

---

Strategies for Handling Missing Values

---

Picking Among Mean, Median, and Mode

The exam frequently poses this as: a continuous feature must be imputed with the least effort and correct results — which strategy? The correct reasoning is to look at the distribution first and choose accordingly, rather than reaching for the mean by default. There is no imputer that inspects the distribution and decides for you.

What every univariate imputation costs

Replacing $k$ missing values with a single constant concentrates mass at that point. The consequences are consistent across mean, median, and mode:

• Variance shrinks — the imputed values have zero spread, so the feature's standard deviation is understated.
• Covariance is diluted — imputed rows carry no relationship to other features, so correlations move toward zero.
• The distribution gains a spike — a histogram after imputation shows an artificial peak at the imputed value, which is worth checking.

These costs are acceptable because the alternative — dropping the rows — is usually worse. But they are the reason a missingness indicator is so often paired with the imputation.

---

Preserving Signal: Missingness Indicator Features

When values are Missing Not at Random (MNAR), the fact that a value is missing is itself a powerful predictive feature. If a data scientist simply imputes the median, the model cannot distinguish between a customer who genuinely had the median value versus a customer who refused to answer.

To solve this, we create a binary missingness indicator: $I(x = \text{null})$.

# Creating explicit missingness indicators in PySpark
from pyspark.sql import functions as F

df_with_indicators = df.withColumn(
    "income_is_missing",
    F.when(F.col("annual_income").isNull(), 1).otherwise(0)
).withColumn(
    "credit_score_is_missing",
    F.when(F.col("credit_score").isNull(), 1).otherwise(0)
)

In scikit-learn, this is automated using SimpleImputer(add_indicator=True) or the dedicated MissingIndicator transformer.

---

Distributed Imputation in PySpark ML (pyspark.ml.feature.Imputer)

Apache Spark ML provides a native, distributed transformer/estimator: pyspark.ml.feature.Imputer. It computes imputation metrics in parallel across the cluster and supports mean, median, and mode strategies.

from pyspark.ml.feature import Imputer

# 1. Instantiate the PySpark Imputer estimator
imputer = Imputer(
    strategy="median",                    # Options: "mean", "median", "mode"
    missingValue=float("nan"),            # Target missing representation
    inputCols=["age", "annual_income", "credit_score"],
    outputCols=["age_imputed", "annual_income_imputed", "credit_score_imputed"]
)

# 2. Fit strictly on the training DataFrame
imputer_model = imputer.fit(train_df)

# 3. Transform both train and test DataFrames using the fitted model
train_imputed = imputer_model.transform(train_df)
test_imputed = imputer_model.transform(test_df)

Scikit-Learn SimpleImputer Example

from sklearn.impute import SimpleImputer
import numpy as np

# Instantiate median imputer with missingness indicator generation
num_imputer = SimpleImputer(strategy="median", add_indicator=True)

# Fit on training split ONLY, transform train and test
X_train_imputed = num_imputer.fit_transform(X_train)
X_test_imputed = num_imputer.transform(X_test)

---

Preventing Data Leakage During Imputation

Data leakage occurs when information from outside the training dataset (such as validation or test splits) is inadvertently used to train or parameterize a machine learning model. This produces unrealistically optimistic offline evaluation metrics that fail to replicate in production.

INCORRECT (Data Leakage!):
[Full Dataset: Train + Test]  ──> Compute Global Median ($65,000) ──> Split Train / Test

CORRECT (Leak-Free ML Pipeline):
[Full Dataset] ──> Split ──┬──> [Train Set] ──> Fit Imputer (Train Median = $62,000)
                           └──> [Test Set]  ──> Transform using Train Median ($62,000)

Rules for Leak-Free Imputation:

1. Split First, Preprocess Second: Always partition raw data into train_df, val_df, and test_df before fitting any imputer.
2. Fit on Train, Transform on Test: Never call .fit() or .fit_transform() on the test dataset. The test set must be treated as unseen real-world production data.
3. Encapsulate in ML Pipelines: Embed transformers inside a pyspark.ml.Pipeline or sklearn.pipeline.Pipeline so that cross-validation folds automatically fit imputers strictly on internal training folds.