Spark SQL对具有相同键的行求和，并附加求和值

Question

假设我有如下的DataFrame:

+----+-----+
|lang|count|
+----+-----+
|  en|    4|
|  en|    5|
|  de|    2|
|  en|    2|
|  nl|    4|
|  nl|    5|
+----+-----+

如何对每种唯一语言的"count"值求和，并将该值附加为新列(因此，不减少行数)?

在我的示例中，这将导致:

+----+-----+----------------+
|lang|count|totalCountInLang|
+----+-----+----------------+
|  en|    4|              11|
|  en|    5|              11|
|  de|    2|               2|
|  en|    2|              11|
|  nl|    4|               9|
|  nl|    5|               9|
+----+-----+----------------+

dataframe是通过map对DStream的操作来构造的。

有什么建议是最好的方法来实现这一点吗?有比使用数据框架更有效的方法吗?

提前感谢!

Answer 1

您可以使用以下方式之一:

• sum over a window:

  import org.apache.spark.sql.expressions.Window
  val df = Seq(
    ("en", 4), ("en", 5), ("de", 2), 
    ("en", 2), ("nl", 4), ("nl", 5)
  ).toDF("lang", "count")
  val w = Window.partitionBy("lang").rowsBetween(
    Window.unboundedPreceding, Window.unboundedFollowing
  )
  df.withColumn("totalCountInLang", sum("count").over(w))
  

• 聚合和连接:

  df.join(df.groupBy("lang").sum(), Seq("lang"))
  

对于小群体，前一种解决方案应该表现得稍微好一些。对于较大的，通常提供后者，可选地与广播功能相结合。