Coercing types and filling nulls without breaking row counts

Nulls are the place pipelines lie quietly. A blanket dropna halves your row count and your dashboard reflects half the data. A blanket fillna(0) plants zeros everywhere and skews every average. The right move is column by column: drop strings, impute numerics.

transform.py
python
def fill_nulls(df: DataFrame) -> DataFrame:
    """Drop string-null rows and fill numeric nulls with column means."""
    null_columns: list[str] = [c for c in df.columns if df.filter(col(c).isNull()).count() > 0]

    string_cols = [
        f.name
        for f in df.schema.fields
        if f.dataType.simpleString() == "string" and f.name in null_columns
    ]
    numeric_cols = [
        f.name
        for f in df.schema.fields
        if f.dataType.simpleString() in ("int", "double") and f.name in null_columns
    ]

    if string_cols:
        df = df.dropna(subset=string_cols)

    if numeric_cols:
        means_row = df.select(*[mean(col(c)).alias(c) for c in numeric_cols]).collect()[0]
        df = df.fillna({c: means_row[c] for c in numeric_cols if means_row[c] is not None})

    return df

The full null handler. Notice it splits string columns from numeric columns and treats them differently. Dropping rows with null carrier is safe. Filling missing arrival delays with the column mean preserves row counts and keeps aggregates honest.

The single .collect() call is doing real work. We compute every numeric column mean in one Spark action, pull the result back as a single Row, and then use it to drive the fillna. That avoids a separate pass per column, which would multiply shuffle cost.

Yes, mean imputation biases. The right answer depends on what null means. Cancelled flights have null arr_delay because no arrival happened. For carrier-level on-time analysis, you might prefer to drop those rows. For schedule reliability analysis, mean is fine. The course shows both choices and when each fits.

Code playground: Pick a fill strategy for your column

Loading practice…

Quiz: Quiz

Loading practice…

AI prompt: Try it: pick a null strategy per column

Loading practice…