PySpark Refactor: One Bad Executor

On-call paged at 02:14 last night: the nightly orders enrichment job ran 3h 40m instead of its usual 22 minutes and one executor OOMed on stage 4. Cluster spend tripled. Volumes only grew about 8% week-over-week. The job lives in jobs/. Reproduce by running tests/test_skew_observability.py. Fix it so the test passes. Be ready to defend the trade-offs of your fix when the interviewer asks what happens with a different data shape.

jobs/enrich_orders.py

# jobs/enrich_orders.py
# Nightly enrichment job. Joins the orders fact with the users dim to
# attach user country, then aggregates daily totals by (order_date,
# country) for the warehouse.

from pyspark.sql import SparkSession, DataFrame
from pyspark.sql import functions as F
from jobs.io import read_orders, read_users, write_enriched

def run(spark: SparkSession) -> int:
    orders: DataFrame = read_orders(spark)   # ~2.4B rows
    users:  DataFrame = read_users(spark)    # ~120M rows, ~9 GB

    enriched = orders.join(users, on="user_id", how="left")

    daily = (
        enriched
        .groupBy("order_date", "country")
        .agg(
            F.count("*").alias("order_count"),
            F.sum("amount_cents").alias("revenue_cents"),
        )
    )

    return write_enriched(daily)

if __name__ == "__main__":
    spark = (
        SparkSession.builder
        .appName("enrich_orders")
        .config("spark.sql.shuffle.partitions", "1000")
        .getOrCreate()
    )
    run(spark)

# jobs/enrich_orders.py
# Nightly enrichment job. Joins the orders fact with the users dim to
# attach user country, then aggregates daily totals by (order_date,
# country) for the warehouse.

from pyspark.sql import SparkSession, DataFrame
from pyspark.sql import functions as F
from jobs.io import read_orders, read_users, write_enriched

def run(spark: SparkSession) -> int:
    orders: DataFrame = read_orders(spark)   # ~2.4B rows
    users:  DataFrame = read_users(spark)    # ~120M rows, ~9 GB

    enriched = orders.join(users, on="user_id", how="left")

    daily = (
        enriched
        .groupBy("order_date", "country")
        .agg(
            F.count("*").alias("order_count"),
            F.sum("amount_cents").alias("revenue_cents"),
        )
    )

    return write_enriched(daily)

if __name__ == "__main__":
    spark = (
        SparkSession.builder
        .appName("enrich_orders")
        .config("spark.sql.shuffle.partitions", "1000")
        .getOrCreate()
    )
    run(spark)

Active Now|Sr. Data Engineer (L5)|||5.5k Attempts|2.1k Solves|