Код-ревью: PySpark-задача отчёта по заказам

orders = spark.read.parquet("s3://dwh/orders/")
users = spark.read.parquet("s3://dwh/users/")
countries = spark.read.parquet("s3://dwh/countries/")

@F.udf("string")
def norm_country(code):
    return code.strip().upper()

df = orders.join(users, "user_id").join(countries, "country_code")
df = df.withColumn("country", norm_country(F.col("country_code")))
df = df.filter(F.col("order_date") == run_date)

print("rows:", df.count())
top = df.groupBy("country").agg(F.sum("amount").alias("revenue"))
rows = df.collect()
top.coalesce(1).write.mode("append").parquet("s3://marts/revenue_by_country/")