title: Pyspark: weighted average by a column
body: <p>For example, I have a dataset like this</p>
  
  <pre><code>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()
  </code></pre>
  
  <p>and I can obtain the customer-region order count matrix by </p>
  
  <pre><code>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()
  </code></pre>
  
  <p>The data would look like this</p>
  
  <p><a href="https://i.stack.imgur.com/4qshi.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/4qshi.png" alt="enter image description here"></a></p>
  
  <p>Now, I want to calculate the weighted average by the <code>overall_count</code>, how can I do it?</p>
  
  <p>The result should be <code>(0.66*3+1*1)/4</code> for region A, and <code>(0.33*3+1*1)/4</code> for region B</p>
  
  <hr>
  
  <p>My thoughts:</p>
  
  <p>It can certainly be achieved through turning the data into python/pandas and then do some calculation, but in what cases should we use Pyspark?</p>
  
  <p>I can get something like </p>
  
  <pre><code>temp_result.agg(sum(col("Region A") * col("overall_count")), sum(col("Region B")*col("overall_count"))).show()
  </code></pre>
  
  <p>but it doesn't feel right, especially if there is many <code>region</code>s to count.</p>