在聚合数据帧中查找平方和分数

Question

假设我有一个简化的DataFrame看起来像这样：

A  B  C
-------
9  6  8 
9  6  8
9  6  7
7  4  2
8  3  2

我想groupBy列A和B，对于每个分区，我想找到唯一值的分数之和C.因此，例如，当按A和B分组时，对于值9和6，我的值将是（2/3）^ 2 +（1/3）^ 2。 3变为分区中有3个值，2来自8的数量，1来自7的数量。这应该针对由groupBy方法创建的每个分区完成。

Answer 1

一种可能的方法是这样的：

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions.{pow, sum}

val w = Window.partitionBy("A", "B")

val tmp = pow(
  $"count" /             // numerator (see below)
  sum($"count").over(w), // denominator
2)

df
  .groupBy($"A", $"B", $"C").count  // Find numerators 
  .withColumn("_tmp", tmp) // Compute fractions
  .groupBy($"A", $"B").agg(sum($"_tmp").as("ssq")) // Compute totals