Your question is Code Optimization. Take a moment with it on the right.
Talk me through your thinking if you like. When you're confident, submit your answer and I'll grade it like a real screen (7/10 or better passes).
Shopee runs this PySpark job nightly to summarize order value per shop. It works on staging but on the full production order volume it runs far longer than the cluster's scheduled window and sometimes fails with driver out-of-memory errors.
from pyspark.sql import functions as F
def summarize_daily_orders(orders_df, shop_metadata_df):
cast_udf = F.udf(lambda x: x * 1.0)
priced = orders_df.withColumn("price_double", cast_udf(F.col("price_cents")))
joined = priced.join(shop_metadata_df, "shop_id")
joined = joined.repartition(1)
grouped = joined.groupBy("shop_id").agg(F.collect_list("price_double").alias("prices"))
rows = grouped.collect()
totals = []
for row in rows:
total = sum(row["prices"])
totals.append((row["shop_id"], total))
order_count = joined.count()
write_count = joined.count()
return totals, order_count, write_count
Optimize this for better time and space complexity on a large, distributed dataset. Explain what you'd change and why, in free text.