How to count a boolean in grouped Spark data frame

Question

I want to count how many of records are true in a column from a grouped Spark dataframe but I don't know how to do that in python. For example, I have a data with a region, salary and IsUnemployed column with IsUnemployed as a Boolean. I want to see how many unemployed people in each region. I know we can do a filter and then groupby but I want to generate two aggregation at the same time as below

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

Adai

Answer 1

Probably the simplest solution is a plain CAST (C style where TRUE -> 1, FALSE -> 0) with SUM:

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

A little bit more universal and idiomatic solution is CASE WHEN with COUNT:

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

but here it is clearly an overkill.

Shrubbery

Does the second approach run slower or does Catalyst optimize both to roughly similar bytecode? – Recount

Answer 2

count_if function

Pyspark 3.5 introduced pyspark.sql.functions.count_if documented as "Returns the number of TRUE values for the col."

So for your example, you could do:

from pyspark.sql import functions as F  
results = (
    data
    .groupby("Region")
    .agg(
        F.avg("Salary").alias("AverageSalary"),
        # new count_if method
        F.count_if("IsUnemployed").alias("CountEmployed"),
        # old casting method still required for getting proportion of true values
        F.avg(F.col("IsUnemployed").cast("integer")).alias("ProportionEmployed"),
    )
)

Southwestward