首页 > 解决方案 > 如何计算和获取 Spark Dataframe 中唯一 ID 的值总和?

问题描述

我有以下数据框,我希望按 id 聚合,并对每个唯一 id 的“值”列求和:

import org.apache.spark.sql.functions._
import spark.implicits._

// some data...
val df = Seq(
  (1, 2),
  (1, 4),
  (1, 1),
  (2, 2),
  (2, 2),
  (3, 2),
  (3, 1),
  (3, 1)
).toDF("id","value")

df.show()

给出以下内容:

+---+-----+
| id|value|
+---+-----+
|  1|    2|
|  1|    4|
|  1|    1|
|  2|    2|
|  2|    2|
|  3|    2|
|  3|    1|
|  3|    1|
+---+-----+

使用 count 函数,我知道我可以计算唯一 ID:

df.select("id").groupBy($"id").count.orderBy($"id".asc).show()

+---+-----+
| id|count|
+---+-----+
|  1|    3|
|  2|    2|
|  3|    3|
+---+-----+

但我也想对每个唯一 ID 的值求和(或求平均值)。所以结果表应该如下:

+---+-----+----------+
| id|count|valueCount|
+---+-----+----------+
|  1|    3|         7|
|  2|    2|         4|
|  3|    3|         4|
+---+-----+----------+

有没有办法以编程方式做到这一点?

标签: apache-sparkdataframe

解决方案


方法是使用聚合函数。Sparks 带有许多预定义的(平均值、总和、计数、第一、收集列表、收集集、最小值、最大值……),因此您始终可以在您的示例中这样做:

df.groupBy("id").agg(
    count("id").as("countOfIds"), 
    sum("id").as("sumOfIds"), 
    avg("id").as("avgOfIds")
).show
+---+----------+--------+--------+
| id|countOfIds|sumOfIds|avgOfIds|
+---+----------+--------+--------+
|  1|         3|       3|     1.0|
|  3|         3|       9|     3.0|
|  2|         2|       4|     2.0|
+---+----------+--------+--------+

您可以通过查看定义为“聚合函数”的函数来查看sql.function 包文档中定义的函数。如果您使用面向 SQL 的语法,所有这些都具有等效的 SQL 语法。


推荐阅读