Adding derived columns before the aggregation

A derived column is a small transform you do once so the rest of the pipeline never has to repeat it. The two we need: total_delay (arrival plus departure delay) and route (origin-dest joined). Both get reused in every window aggregation downstream.

transform.py
python
def add_derived_columns(df: DataFrame) -> DataFrame:
    """Add total_delay and route columns used by the downstream aggregations."""
    df = df.withColumn("total_delay", col("arr_delay") + col("dep_delay"))
    df = df.withColumn("route", concat_ws("-", col("origin"), col("dest")))
    return df

Two `withColumn` calls. The total_delay sum is the simple addition. The route column uses `concat_ws` to join origin and dest with a hyphen, which gives us a stable string key for partitioning later.

Computing total_delay inside every window function would mean Catalyst pushes the same expression into multiple plan branches. Defining it once up front lets the optimizer evaluate it once and reuse the result across every downstream aggregation.

Why concat_ws instead of concat? concat returns null if any input is null, which would silently drop rows from your route column. concat_ws skips nulls and joins what is present. Defensive default for any string composition that uses real data.

Pushing derived columns into the source means the source has to know what every consumer needs. That couples the producer to every downstream job. Keep derivations in the transform layer where they live next to the queries that use them.

Quiz: Quiz

Loading practice…

AI prompt: Try it: spot the missing derived columns

Loading practice…