scala - 在 Spark 中替代 groupBy

Question

我有一个带有下一个信息的 Dataframe df:

id   json_data
1    {a: "1", b: "2"} 
1    {a: "1", b: "3"} 
1    {a: "1", b: "4"}
2    {a: "1", b: "2"} 
2    {a: "1", b: "6"}

我需要下一个最终结果:

id   json_data
1    [{a: "1", b: "2"},{a: "1", b: "3"},{a: "1", b: "4"}]
2    [{a: "1", b: "2"},{a: "1", b: "6"}]

我尝试了两种不同的方法，分别使用 Window 函数和 groupBy。通过这两种方法，我都得到了想要的结果。

1º 方法:

var user_window = Window.partitionBy("id").orderBy("id")
val df2 = df.withColumn("json_data",
    collect_list($"json_data").over(user_window))
          .withColumn("rank", row_number().over(user_window))
          .where("rank = 1")

2º 方法:

val df2 = df.groupBy(df("id")).agg(collect_list($"json_data").as("json_data"))

通过这两种方法，我获得了相同的性能。但是阅读有关 Spark 的文档，似乎这两种方法都不高效，因为具有相同键的行将需要穿过集群(洗牌)才能在一起。我展示了一个小例子，因为在生产中我有大量数据。做组或使用 Window 函数需要很长时间。

为了做到这一点，有什么替代方案吗？

Answer 1

我的建议是使用 reduceByKey。

这样，如果您的键是 id 并且您的值(在开始时)是列表中的 json_data ，则执行 reduceByKey 以及串联功能，各种 json_data 包装列表将为您提供更好的性能。

简而言之，首先使用 reduceByKey 在分区内执行“groupBy”，然后才开始数据混洗。

阅读 groupByKey 和 reduceByKey 的性能差异的好地方是 here (部分 6b)。

在 pyspark 中它看起来像这样:

rdd = df.rdd
rdd = rdd.map(lambda row: (row['id'], [row['json_data']]))
rdd = rdd.reduceByKey(lambda a, b: a + b)