apache-spark - Spark DataFrame 基于条件的列总和

Question

我想计算值的部分，只有两个分区(其中 type == red 和 where type != red)

ID     | type        | value
-----------------------------
  1    |    red      |  10
  2    |    blue     |  20
  3    |    yellow   |  30

结果应该是:

ID     | type        | value | portion
-----------------------------
  1    |    red      |  10   | 1
  2    |    blue     |  20   |0.4
  3    |    yellow   |  30   |0.6

spark中的普通窗口函数仅支持按整列分区，但我需要“蓝色”和“黄色”，一起识别为“非红色”类型。

有什么想法吗？

Answer 1

首先添加一列 is_red 以便更轻松地区分两组。然后，您可以groupBy这个新列并分别获取两个组的总和。

要获得分数(部分)，只需将每行的值除以正确的总和，同时考虑类型是否为红色。这部分可以在 Spark 中使用 when 和 otherwise 来完成。

下面是执行此操作的 Scala 代码。由于使用 groupBy 时无法保证结果的顺序，因此存在 sortBy。通过排序，下面的 sum1 将包含所有非红色类型的总和，而 sum2 是红色类型的总和。

val sum1 :: sum2 :: _ = df.withColumn("is_red", $"type" === lit("red"))
  .groupBy($"is_red")
  .agg(sum($"value"))
  .collect()
  .map(row => (row.getAs[Boolean](0), row.getAs[Long](1)))
  .toList
  .sortBy(_._1)
  .map(_._2)

val df2 = df.withColumn("portion", when($"is_red", $"value"/lit(sum2)).otherwise($"value"/lit(sum1)))

可以使用 drop 删除额外的 is_red 列。