Refactoring for Scalability

JOHN LEONARD's data platform team runs a weekly PySpark job that reports total placement fees earned by each recruitment consultant, built from raw booking records sitting in the data lake. As booking volume has grown, the job now regularly times out, and it has crashed the driver outright a few times right before the report was due.

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf, collect_list
from pyspark.sql.types import DoubleType

spark = SparkSession.builder.appName("consultant_fee_report").getOrCreate()

bookings_df = spark.read.parquet("s3://john-leonard-dw/bookings/")

@udf(returnType=DoubleType())
def to_gbp(amount, fx_rate):
    return amount * fx_rate

bookings_df = bookings_df.withColumn("fee_gbp", to_gbp(bookings_df.fee_amount, bookings_df.fx_rate))

all_rows = bookings_df.collect()
filtered = [r for r in all_rows if r.status == "PLACED"]

placed_df = spark.createDataFrame(filtered, bookings_df.schema)

grouped = placed_df.groupBy("consultant_id").agg(collect_list("fee_gbp").alias("fees"))

results = []
for row in grouped.collect():
    total = 0.0
    for fee in row.fees:
        total += fee
    results.append((row.consultant_id, total))

report_df = spark.createDataFrame(results, ["consultant_id", "total_fee_gbp"])
report_df.repartition(1).write.mode("overwrite").parquet("s3://john-leonard-dw/reports/weekly_fees/")

Explain what you would change in this job to make it scale to a much larger booking volume, and why each change matters.