Aggregating columns conditionally with pyspark? - apache-spark

Question

I have following dataset. I want to group all variables and split the data based on the conditions below.
However, I am getting error when I tried the code below.
CUST_ID NAME GENDER AGE
id_01 MONEY F 43
id_02 BAKER F 32
id_03 VOICE M 31
id_04 TIME M 56
id_05 TIME F 24
id_06 TALENT F 28
id_07 ISLAND F 21
id_08 ISLAND F 27
id_09 TUME F 24
id_10 TIME F 75
id_11 SKY M 35
id_12 VOICE M 70
from pyspark.sql.functions import *
df.groupBy("CUST_ID", "NAME", "GENDER", "AGE").agg(
CUST_ID.count AS TOTAL
SUM(WHEN ((AGE >= 18 AND AGE <= 34) AND GENDER = 'M') THEN COUNT(CUST_ID) ELSE 0 END AS "M18-34")
SUM(WHEN ((AGE >= 18 AND AGE <= 34) AND GENDER = 'F') THEN COUNT(CUST_ID) ELSE 0 END AS "F18-34")
SUM(WHEN ((AGE >= 18 AND AGE <= 34 THEN COUNT(CUST_ID) ELSE 0 END AS "18-34")
SUM(WHEN ((AGE >= 25 AND AGE <= 54 THEN COUNT(CUST_ID) ELSE 0 END AS "25-54")
SUM(WHEN ((AGE >= 25 AND AGE <= 54) AND GENDER = 'F') THEN COUNT(CUST_ID) ELSE 0 END AS "F25-54")
SUM(WHEN ((AGE >= 25 AND AGE <= 54) AND GENDER = 'M') THEN COUNT(CUST_ID) ELSE 0 END AS "M25-54")
)
I would appreciate your help/suggestions
Thanks in advance

Answer 1

Your code is neither valid pyspark nor valid Spark SQL. There are so many syntax problems. I attempted to fix them below, not sure if that's what you wanted. If you have so many SQL-like statements, it's better to use Spark SQL directly rather than the pyspark API:
df.createOrReplaceTempView('df')
result = spark.sql("""
SELECT NAME,
COUNT(CUST_ID) AS TOTAL,
SUM(CASE WHEN ((AGE >= 18 AND AGE <= 34) AND GENDER = 'M') THEN 1 ELSE 0 END) AS `M18-34`,
SUM(CASE WHEN ((AGE >= 18 AND AGE <= 34) AND GENDER = 'F') THEN 1 ELSE 0 END) AS `F18-34`,
SUM(CASE WHEN (AGE >= 18 AND AGE <= 34) THEN 1 ELSE 0 END) AS `18-34`,
SUM(CASE WHEN (AGE >= 25 AND AGE <= 54) THEN 1 ELSE 0 END) AS `25-54`,
SUM(CASE WHEN ((AGE >= 25 AND AGE <= 54) AND GENDER = 'F') THEN 1 ELSE 0 END) AS `F25-54`,
SUM(CASE WHEN ((AGE >= 25 AND AGE <= 54) AND GENDER = 'M') THEN 1 ELSE 0 END) AS `M25-54`
FROM df
GROUP BY NAME
""")
result.show()
+------+-----+------+------+-----+-----+------+------+
| NAME|TOTAL|M18-34|F18-34|18-34|25-54|F25-54|M25-54|
+------+-----+------+------+-----+-----+------+------+
|ISLAND| 2| 0| 2| 2| 1| 1| 0|
| MONEY| 1| 0| 0| 0| 1| 1| 0|
| TIME| 3| 0| 1| 1| 0| 0| 0|
| VOICE| 2| 1| 0| 1| 1| 0| 1|
| TUME| 1| 0| 1| 1| 0| 0| 0|
| BAKER| 1| 0| 1| 1| 1| 1| 0|
|TALENT| 1| 0| 1| 1| 1| 1| 0|
| SKY| 1| 0| 0| 0| 1| 0| 1|
+------+-----+------+------+-----+-----+------+------+
If you want a pyspark solution, here's an example of how to do it for the first column. You can work out the rest straightforwardly.
import pyspark.sql.functions as F
result = df.groupBy('Name').agg(
F.count('CUST_ID').alias('TOTAL'),
F.count(F.when(F.expr("(AGE >= 18 AND AGE <= 34) AND GENDER = 'M'"), 1)).alias("M18-34")
)
result.show()
+------+-----+------+
| Name|TOTAL|M18-34|
+------+-----+------+
|ISLAND| 2| 0|
| MONEY| 1| 0|
| TIME| 3| 0|
| VOICE| 2| 1|
| TUME| 1| 0|
| BAKER| 1| 0|
|TALENT| 1| 0|
| SKY| 1| 0|
+------+-----+------+