
Image by Author | Canva
Did you know that 402.7 million terabytes of data are created each day? This amount of data that has been collected needs to be aggregated to find hidden insights or discover trends. PySpark is the go-to tool for that. We will discover how you can use basic or advanced aggregations using actual interview datasets! Let’s get started!
Basic Aggregation
In this section, we will explore basic aggregation, such as mean(), min(), max(), count(), and average (). To do that, we will use the amazon_purchases dataset from the Revenue Over Time question. Amazon has asked this question, and you can reach out here.
Revenue Over Time
Find the 3-month rolling average of total revenue from purchases given a table with users, their purchase amount, and date purchased. Do not include returns which are represented by negative purchase values. Output the year-month (YYYY-MM) and 3-month rolling average of revenue, sorted from earliest month to latest month. A 3-month rolling average is defined by calculating the average total revenue from all user purchases for the current month and previous two months. The first two months will not be a true 3-month rolling average since we are not given data from last year. Assume each month has at least one purchase.
Here is the preview of this dataset.

This dataset contains three columns: user_id, created_at, and purchase_amt, and it’s pretty straightforward to discover these basic aggregation methods.
min() and max()
Using those aggregation functions, we can perform all kinds of mathematical operations. So, when it comes to calculations, finding the biggest or smallest number comes to mind first. So, let’s see the minimum and maximum purchases. Here is the code.
from pyspark.sql.functions import min, max
# Calculate the minimum purchase amount
min_purchase = amazon_purchases.select(min("purchase_amt").alias("Minimum Purchase"))
# Calculate the maximum purchase amount
max_purchase = amazon_purchases.select(max("purchase_amt").alias("Maximum Purchase"))
# Show the results
min_purchase.show()
max_purchase.show()
Here is the output.

count()
Let’s find the number of rows in this dataset. This is a pretty straightforward method for understanding its size.
from pyspark.sql.functions import count
# Count the total number of rows in the dataset
row_count = amazon_purchases.select(count("*").alias("Total Rows"))
# Show the result
row_count.show()
Here is the output.

Good, as you can see, we have found total rows.
sum()
Now, let’s calculate the total sales. To do that, we will use the sum() function to add up all the values in the purchase_amt column. Here is the code.
from pyspark.sql.functions import sum
# Calculate the total sales
total_sales = amazon_purchases.select(sum("purchase_amt").alias("Total Sales"))
# Show the result
total_sales.show()
Here is the output.

avg()
Good. Do you wonder what is the average purchase amount? For this, let’s use the avg() function to compute the mean of purchase_amt. Here is the code you can use.
from pyspark.sql.functions import avg
# Calculate the average purchase amount
average_purchase = amazon_purchases.select(avg("purchase_amt").alias("Average Purchase"))
# Show the result
average_purchase.show()
Here is the output.

Grouped Operations
Now that you understand data aggregation, let’s move forward with more advanced operations.
Grouped aggregation is a good way to summarize data by grouping it. Instead of looking at the entire dataset, this method finds patterns or trends in smaller groups of data.
By now, you are familiar with this dataset, so let’s use the same dataset and find out its total monthly revenue by grouping purchases based on their month and year. Here is the code.
import pandas as pd
import numpy as np
from datetime import datetime
from pyspark.sql import functions as F
from pyspark.sql.window import Window
pd.options.display.float_format = "{:,.2f}".format
df = amazon_purchases.filter(amazon_purchases['purchase_amt'] > 0)
df = df.withColumn('month_year', F.date_format(F.col('created_at'), 'yyyy-MM'))
df1 = df.groupby('month_year').agg(F.sum('purchase_amt').alias('monthly_revenue')).orderBy('month_year')
df1.show()
Here is the output.

What did we do here? Here, the groupby() function groups the data by month and year. So this will allow us to calculate the total revenue for each month separately. Simply put, we track monthly income over time to see the dataset’s financial performance.
Window Aggregation
Good. Now, we understand the easier and more advanced usage of aggregation functions. So, let’s look at window aggregation, the most advanced technique we will explain in this article.
Window functions let you do calculations on a set “window” of rows without combining them into a single summary. Since the relationship between rows is important for operations like rolling averages, ranks, and cumulative sums, these are the best types of tables for you.
To do that, let’s look at Amazon’s question again. Finally, we will solve it completely.
In this question, Amazon simply asks us to compute the average revenue for the current month and the previous two months while excluding returns (negative values). The first two months won’t have a complete rolling average since we lack data from the previous year. In the previous section, we found the total monthly revenue by grouping purchases based on their month and year, forming this solution’s first part. We will then figure out the three-month rolling average of the monthly income. This involves setting up a three-month rolling window and computing the average income for the current month and the two months before. To see, check the code below:
df1 = df1.withColumn('month_year', F.to_date(F.col('month_year'), 'yyyy-MM'))
df1 = df1.withColumn(
'rolling_mean',
F.avg('monthly_revenue').over(Window.orderBy('month_year').rowsBetween(-2, 0))
)
result = df1.select('month_year', 'rolling_mean').toPandas()
result
So, our entire code will become this if we combine it with the previous section:
import pandas as pd
import numpy as np
from datetime import datetime
from pyspark.sql import functions as F
from pyspark.sql.window import Window
pd.options.display.float_format = "{:,.2f}".format
df = amazon_purchases.filter(amazon_purchases['purchase_amt'] > 0)
df = df.withColumn('month_year', F.date_format(F.col('created_at'), 'yyyy-MM'))
df1 = df.groupby('month_year').agg(F.sum('purchase_amt').alias('monthly_revenue')).orderBy('month_year')
df1 = df1.withColumn('month_year', F.to_date(F.col('month_year'), 'yyyy-MM'))
df1 = df1.withColumn('rolling_mean', F.avg('monthly_revenue').over(Window.orderBy('month_year').rowsBetween(-2, 0)))
result = df1.select('month_year', 'rolling_mean').toPandas()
result
Here is the output.

If you want to know more about PySpark, check out this one: What is PySpark?
Common Pitfalls to Avoid in Data Aggregation
Now, we have discovered Data aggregation at each level. But before doing that, let’s look at common pitfalls to avoid to make our codes even better in the future.
- Null and Missing Values: Null values and Missing Values can distort results in data aggregation functions like sum(), mean(), or avg(), so be careful checking these first.
- Skipping Data Filtering: Any extra data, such as returns (negative values), can universally affect results. Filters can be applied to aggregate only relevant data points.
- Misusing GroupBy: Too many or irrelevant fields in a grouping result in bloated or wrong results. Concentrate on significant clusters that support your analysis objectives.
- Failure to Optimize Performance: Poor resource or partition allocation also degrades performance.So, be careful when optimizing partitions and memory usage, especially in big datasets.
Conclusion
Data aggregation is a cornerstone of practical data analysis. As the amount of data collected has dramatically increased daily, knowing these techniques, especially by using the go-to tool for industries like Pyspark, is crucial for your career.
One important thing is to learn these techniques by using real-life interview questions and data projects, which can be counted as an experience in the future!
