CODEWITHSUNDEEP
pythonapache-sparkpyspark
martina.physics
martina.physics

Reputation: 9804

PySpark count values by condition

I have a DataFrame, a snippet here:

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

basically a string field named f and either a 1 or a 0 for second element (is_fav).

What I need to do is grouping on the first field and counting the occurrences of 1s and 0s. I was hoping to do something like

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)

It does not work properly, I get in both cases the same result which amounts to the count for the items in the group, so the filter (whether it is a 1 or a 0) seems to be ignored. Does this depend on how count works?

Upvotes: 8

Views: 27866

Answers (1)

zero323
zero323

Reputation: 330423

There is no filter here. Both col("is_fav") == 1 and col("is_fav") == 0) are just boolean expressions and count doesn't really care about their value as long as it is defined.

There are many ways you can solve this for example by using simple 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")
)

or making ignored values undefined (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)

Upvotes: 16

Related Questions