You can use the following methods to sum the values in a column of a PySpark DataFrame that meet a condition:
Method 1: Sum Values that Meet One Condition
from pyspark.sql.functions import sum
#sum values in points column for rows where team column is 'B'
df.filter(df.team=='B').agg(sum('points')).collect()[0][0]
Method 2: Sum Values that Meet Multiple Conditions
from pyspark.sql.functions import sum
#sum values in points column for rows where team is 'B' and position is 'Guard'
df.filter((df.team=='B') & (df.position=='Guard')).agg(sum('points')).collect()[0][0]
Method 3: Sum Values that Meet One of Several Conditions
from pyspark.sql.functions import sum
#sum values in points column for rows where team is 'B' or position is 'Guard'
df.filter((df.team=='B') | (df.position=='Guard')).agg(sum('points')).collect()[0][0]
The following examples show how to use each method in practice with the following PySpark DataFrame that contains information about various basketball players:
from pyspark.sql import SparkSession
spark = SparkSession.builder.getOrCreate()
#define data
data = [['A', 'Guard', 11],
['A', 'Guard', 8],
['A', 'Forward', 22],
['A', 'Forward', 22],
['B', 'Guard', 14],
['B', 'Guard', 14],
['B', 'Forward', 13],
['B', 'Forward', 7]]
#define column names
columns = ['team', 'position', 'points']
#create dataframe using data and column names
df = spark.createDataFrame(data, columns)
#view dataframe
df.show()
+----+--------+------+
|team|position|points|
+----+--------+------+
| A| Guard| 11|
| A| Guard| 8|
| A| Forward| 22|
| A| Forward| 22|
| B| Guard| 14|
| B| Guard| 14|
| B| Forward| 13|
| B| Forward| 7|
+----+--------+------+
Example 1: Sum Values that Meet One Condition
We can use the following syntax to sum the values in the points column where the corresponding value in the team column is equal to B:
from pyspark.sql.functions import sum
#sum values in points column for rows where team column is 'B'
df.filter(df.team=='B').agg(sum('points')).collect()[0][0]
48
We can see that the sum of values in the points column for players on team B is 48.
Example 2: Sum Values that Meet Multiple Conditions
We can use the following syntax to sum the values in the points column where the corresponding value in the team column is equal to B and the value in the position column is equal to Guard:
from pyspark.sql.functions import sum
#sum values in points column for rows where team is 'B' and position is 'Guard'
df.filter((df.team=='B') & (df.position=='Guard')).agg(sum('points')).collect()[0][0]
28
We can see that the sum the values in the points column where the corresponding value in the team column is equal to B and the value in the position column is equal to Guard is 28.
Example 3: Sum Values that Meet One of Several Conditions
We can use the following syntax to sum the values in the points column where the corresponding value in the team column is equal to B or the value in the position column is equal to Guard:
from pyspark.sql.functions import sum
#sum values in points column for rows where team is 'B' or position is 'Guard'
df.filter((df.team=='B') | (df.position=='Guard')).agg(sum('points')).collect()[0][0]
67
We can see that the sum the values in the points column where the corresponding value in the team column is equal to B or the value in the position column is equal to Guard is 67.
Additional Resources
The following tutorials explain how to perform other common tasks in PySpark:
How to Calculate a Cumulative Sum in PySpark
How to Calculate Sum by Group in PySpark
How to Sum Multiple Columns in PySpark
Can I obtain the same results using when() function instead of filter()??
I’m using this example using when() but I don’t understand it. Can you explain me??
(I must use when() for a personal project)
df.agg(sum(
when(df.team==’B’,’points’).otherwise(1.))
.alias(“XX”)).show()
I obtain 4 points, I don’t understand this result
Hi Pablo…Yes, you can use the `when()` function instead of `filter()` to conditionally sum a column in PySpark. However, your code contains a mistake that affects the result.
### **Why is the result 4?**
Your current code:
“`python
df.agg(sum(
when(df.team == ‘B’, ‘points’).otherwise(1.)
).alias(“XX”)).show()
“`
– **`when(df.team == ‘B’, ‘points’)`** → This part should return the value of the `”points”` column **only** when `team == ‘B’`.
– **`.otherwise(1.)`** → This means that **for all other rows**, it assigns the value `1.0` instead of ignoring them.
– The `sum()` function then sums **both the `”points”` values for team ‘B’ and `1.0` for all other rows**.
### **How to Fix It?**
If you want to sum `”points”` **only** where `team == ‘B’`, then you should use:
“`python
from pyspark.sql.functions import sum, when
df.agg(sum(when(df.team == ‘B’, df.points).otherwise(0)).alias(“XX”)).show()
“`
### **Explanation of Fix:**
– `when(df.team == ‘B’, df.points)` → If `team == ‘B’`, take `”points”`, else `None` (null).
– `.otherwise(0)` → Ensures that rows not meeting the condition contribute `0` instead of `1.0`.
– `sum(…)` → Correctly sums only the `”points”` for `team == ‘B’`.
### **Alternative Using `filter()`**
If we used `filter()`, the equivalent code would be:
“`python
df.filter(df.team == ‘B’).agg(sum(df.points).alias(“XX”)).show()
“`
This filters the DataFrame first and then sums `”points”`, which is more intuitive.
Thanks so much for your reply @James
Now is OK 🙂
df.agg(sum(
when(df.team==’B’,df.points).otherwise(0))
.alias(“XX”)).show()
print(‘puntos B >> ‘,df.filter(df.team==’B’).agg(sum(‘points’)).collect()[0][0])
You are very welcome Pablo! Keep us posted on your progress!