如何计算分组的 Spark 数据框中的布尔值

Question

我想计算分组的 Spark 数据框中的一列中有多少条记录为真，但我不知道如何在 python 中执行此操作。例如，我有一个包含 region、salary 和 IsUnemployed 列的数据，其中 IsUnemployed 作为布尔值。我想看看每个地区有多少失业人员。我知道我们可以做一个 filter 然后 groupby 但我想同时生成两个聚合如下

from pyspark.sql import functions as F  
data.groupby("Region").agg(F.avg("Salary"), F.count("IsUnemployed")) 

Answer 1

可能最简单的解决方案是 CAST（C 风格，其中 TRUE -> 1，FALSE -> 0），SUM:

(data
    .groupby("Region")
    .agg(F.avg("Salary"), F.sum(F.col("IsUnemployed").cast("long"))))

更通用和惯用的解决方案是 CASE WHEN 和 COUNT:

(data
    .groupby("Region")
    .agg(
        F.avg("Salary"),
        F.count(F.when(F.col("IsUnemployed"), F.col("IsUnemployed")))))

但这里显然是矫枉过正了。