PySpark按条件计算值

时间:2016-03-17 21:46:12

标签: python apache-spark pyspark

我有一个DataFrame,这里有一个片段:

[['u1', 1], ['u2', 0]]

基本上是一个名为f的字符串字段,第二个元素(is_fav)为1或0。

我需要做的是对第一个字段进行分组并计算1和0的出现次数。我希望做一些像

这样的事情
num_fav = count((col("is_fav") == 1)).alias("num_fav")

num_nonfav = count((col("is_fav") == 0)).alias("num_nonfav")

df.groupBy("f").agg(num_fav, num_nonfav)

它无法正常工作,我在两种情况下都得到相同的结果,这相当于组中项目的计数,因此过滤器(无论是1还是0)似乎被忽略。这取决于count的工作原理吗?

1 个答案:

答案 0 :(得分:13)

这里没有过滤器。 col("is_fav") == 1col("is_fav") == 0)都只是布尔表达式,count只要定义它们就不关心它们的值。

有很多方法可以解决这个问题,例如使用简单的sum

from pyspark.sql.functions import sum, abs

gpd = df.groupBy("f")
gpd.agg(
    sum("is_fav").alias("fv"),
    (count("is_fav") - sum("is_fav")).alias("nfv")
)

或使忽略值未定义(a.k.a NULL):

exprs = [
    count(when(col("is_fav") == x, True)).alias(c)
    for (x, c) in [(1, "fv"), (0, "nfv")]
]
gpd.agg(*exprs)