Store count as variable and use it for calculations in PySpark - sql

Question

I have dataframe df1:
+------+-----------+----------+----------+-----+
| sid|acc_term_id|first_name| last_name|major|
+------+-----------+----------+----------+-----+
|106454| 2014B| Doris| Marshall| BIO|
|106685| 2015A| Sara|Richardson| CHM|
|106971| 2015B| Rose| Butler| CHM|
|107298| 2015B| Kayla| Barnes| CSC|
|107555| 2016A| Carolyn| Ford| PHY|
|107624| 2016B| Marie| Webb| BIO|
I want to store the count of sid from this dataframe
c_value = current.agg({"sid": "count"}).collect()[0][0]
and use it for creating a prop column as shown in code below:
c_value = current.agg({"sid": "count"}).collect()[0][0]
stud_major = (
current
.groupBy('major')
.agg(
expr('COUNT(*) AS n_students')
)
.select('major', 'n_students', expr('ROUND(n_students/c_value, 4) AS prop'),
)
)
stud_major.show(16)
When I run the code I get error:
cannot resolve '`c_value`' given input columns: [major, n_students]; line 1 pos 17;
If I put numeric value 2055 instead of c_value everything ok like below.
+
-----+----------+------+
|major|n_students| prop|
+-----+----------+------+
| MTH| 320|0.1557|
| CHM| 405|0.1971|
| CSC| 508|0.2472|
| BIO| 615|0.2993|
| PHY| 207|0.1007|
+-----+----------+------+
Probably there are other ways to calculate but I need by storing count as variable.
Any ideas?

Answer 1

In jupyter Use pandas agg
j=df.agg({'sid':'count'})
df.groupby("major")['sid'].agg(n_students=(lambda x: x.count()), prop=(lambda x: x.count()/j))
major n_students prop
0 BIO 2 0.333333
1 CHM 2 0.333333
2 CSC 1 0.166667
3 PHY 1 0.166667
and pyspark
from pyspark.sql.functions import *
df.groupby('major').agg(count('sid').alias('n_students')).withColumn('prop', round((col('n_students')/c_value),2)).show()
Alternatively You could
c_value = df.agg({"sid": "count"}).collect()[0][0]
df.groupBy('major').agg(expr('COUNT(*) AS n_students')).selectExpr('major',"n_students", f"ROUND(n_students/{c_value},2) AS prop").show()
+-----+----------+----+
|major|n_students|prop|
+-----+----------+----+
| BIO| 2|0.33|
| CHM| 2|0.33|
| CSC| 1|0.17|
| PHY| 1|0.17|
+-----+----------+----+