基于条件的列的火花数据帧总和

Question

我想计算值的一部分，只有两个分区（where 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中普通的window函数只支持partitionby整列，但我需要"blue"和"yellow"，一起识别为"non-red"类型。

有什么想法吗？

Answer 1

首先添加一列is_red以便更容易区分两组。然后你可以groupBy这个新列并分别得到两组的总和。

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

下面是执行此操作的 Scala 代码。有一个 sortBy，因为当使用 groupBy 时，结果的顺序是无法保证的。通过排序，下面的 sum1 将包含所有 non-red 类型的总和，而 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 列。

Answer 2

受 Shaido 的启发，我使用了一个额外的列 is_red 和 spark window 函数。但是我不确定哪个性能更好。

 df.withColumn("is_red", when(col("type").equalTo("Red"), "Red")
   .otherwise("not Red")
   .withColumn("portion", col("value")/sum("value)
   .over(Window.partitionBy(col"is_Red")))
   .drop(is_Red)

基于条件的列的火花数据帧总和

spark dataframe sum of column based on condition

dataframe

apache-spark

apache-spark-sql

spark-dataframe