将列添加到Spark / scala中具有特定条件的数据框

Question

我有一个表格的数据框（如下所示），其中每个id都有一个相应的元素数量。该桶需要3个低值，中等和高值。

+---+------+-----+
| id|bucket|count|
+---+------+-----+
|id1|   low|  2  |
|id1|   med|  3  |
|id1|  high|  4  |
|id2|   low|  1  |
|id2|   med|  4  |
|id3|   low|  7  |
|id3|  high|  1  |
|id4|   med|  2  |
|id4|  high|  1  |
+---+------+-----+ 

我想要的输出如下

+---+-----+-----+-----+
|id | low | med | high|
+---+-----+-----+-----+
|id1|    2|  3  |  4  |
|id2|    1|  4  |  0  |
|id3|    7|  0  |  1  |
|id4|    0|  2  |  1  |
+---+-----+-----+-----+

如果该特定id的存储桶没有条目，则输出中的计数默认为零。

我是新来的火花，无法弄清楚获得此结果的查询。最终的架构是固定的。

Answer 1

您可以使用pivot功能获取输出

df.groupBy("id").pivot("bucket").agg(sum("count"))
  .na.fill(0)
  .show(false)

输出：

+---+----+---+---+
|id |high|low|med|
+---+----+---+---+
|id3|1   |7  |0  |
|id1|4   |2  |3  |
|id2|0   |1  |4  |
|id4|1   |0  |2  |
+---+----+---+---+

希望这有帮助！