scala - 如何在 Spark 中将数据帧设为 "dense"

Question

我有一个数据框，如下所示:

item_id  week_id  sale amount
1           1         10
1           2         12
1           3         15
2           1         4
2           2         7
2           3         9

我想将此数据框转换为新的数据框，如下所示:

item_id   week_1     week_2     week_3
1          10          12         15
2          4            7          9

这可以在 R 中轻松完成，但我不知道如何使用 Spark API 和 Scala 来完成。

Answer 1

您可以使用groupBy.pivot，然后聚合sale_amount列，在这种情况下，您可以从项目和周的每个组合ID中获取第一个值(如果有)每个组合中不超过一行:

df.groupBy("item_id").pivot("week_id").agg(first("sale_amount")).show
+-------+---+---+---+
|item_id|  1|  2|  3|
+-------+---+---+---+
|      1| 10| 12| 15|
|      2|  4|  7|  9|
+-------+---+---+---+

如果 item_id 和 week_id 的每个组合都有不止一行，您可以使用其他聚合函数，例如 sum:

df.groupBy("item_id").pivot("week_id").agg(sum("sale_amount")).show
+-------+---+---+---+
|item_id|  1|  2|  3|
+-------+---+---+---+
|      1| 10| 12| 15|
|      2|  4|  7|  9|
+-------+---+---+---+

要获得正确的列名称，您可以在旋转之前转换 week_id 列:

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

(df.withColumn("week_id", concat(lit("week_"), df("week_id"))).
    groupBy("item_id").pivot("week_id").agg(first("sale_amount")).show)

+-------+------+------+------+
|item_id|week_1|week_2|week_3|
+-------+------+------+------+
|      1|    10|    12|    15|
|      2|     4|     7|     9|
+-------+------+------+------+