Pyspark: средневзвешенная по столбцу

Question

Например, у меня есть такой набор данных

test = spark.createDataFrame([
    (0, 1, 5, "2018-06-03", "Region A"),
    (1, 1, 2, "2018-06-04", "Region B"),
    (2, 2, 1, "2018-06-03", "Region B"),
    (3, 3, 1, "2018-06-01", "Region A"),
    (3, 1, 3, "2018-06-05", "Region A"),
])\
  .toDF("orderid", "customerid", "price", "transactiondate", "location")
test.show()

и я могу получить матрицу подсчета заказов по регионам клиентов

overall_stat = test.groupBy("customerid").agg(count("orderid"))\
  .withColumnRenamed("count(orderid)", "overall_count")
temp_result = test.groupBy("customerid").pivot("location").agg(count("orderid")).na.fill(0).join(overall_stat, ["customerid"])

for field in temp_result.schema.fields:
    if str(field.name) not in ['customerid', "overall_count", "overall_amount"]:
        name = str(field.name)
        temp_result = temp_result.withColumn(name, col(name)/col("overall_count"))
temp_result.show()

Данные будут выглядеть так

Теперь я хочу вычислить средневзвешенное значение по overall_count, как я могу это сделать?

Результат должен быть (0.66*3+1*1)/4 для региона A и (0.33*3+1*1)/4 для региона B

---

Мои мысли:

Конечно, это можно сделать, превратив данные в python / pandas, а затем выполнить некоторые вычисления, но в каких случаях мы должны использовать Pyspark?

Я могу получить что-то вроде

temp_result.agg(sum(col("Region A") * col("overall_count")), sum(col("Region B")*col("overall_count"))).show()

но это не так, особенно если учесть много region с.