Image by Author | Canva
According to a report by Statista, the total amount of data created, consumed, and stored is predicted to reach 149 zettabytes in 2024, which is expected to reach 394 in 2028.
Image Source: Statista
This massive amount of data needs to be processed, right? Mathematic operations on the data to answer questions that add meaning to the data are done using window functions.
In this article, we’ll use real-life examples; at the end, you can solve Visa interview questions. So let’s get started!
What Are PySpark Window Functions?
PySpark window functions allow computations across a set of rows somewhat connected to the current row without collapsing the rows into a single output row.
While window functions preserve the structure of the original, allowing a small step back so that complex insight and richer insights may be drawn, classic aggregate functions aggregate a dataset, reducing it to a more informed version of the original.
But let’s first look at PySpark window function types and then the practical examples.
PySpark Ranking Functions
In this section, we will use three different ranking functions to allow you to sort datasets based on your requests. Let’s review useful cases now using the following datasets.
Our first dataset is transaction_records:
The second dataset is stores.
Here is the interview question on high-density areas:
High-Density Areas
Identify the top 3 areas with the highest customer density. Customer density = (total number of unique customers in the area / area size). Your output should include the area name and its calculated customer density, and ties will be ranked the same.
Link to this question: https://platform.stratascratch.com/coding/10544-high-density-areas
At the end of the article, we will solve this interview question. Let’s first answer an easier question using these datasets so we can build up to solving this question easily.
ROW_NUMBER() Function
The ROW_NUMBER() function provides each row in a partition with a unique sequential number. Unlike other ranking functions, there are no ties, and each row receives a unique rank.
For example, how would we find each region’s five most significant transactions?
To resolve this, we take the transaction_records dataset, join the stores table on the column store_id to get area_name related to transactions, partition by area_name and order by transaction_amount in descending order, and then apply the ROW_NUMBER() function to be ranked.
from pyspark.sql.window import Window
from pyspark.sql.functions import col, row_number
transactions_with_area = transaction_records.join(
stores, transaction_records.store_id == stores.store_id
).select(
"transaction_id", "customer_id", "transaction_amount", "transaction_date", "area_name"
)
window_spec = Window.partitionBy("area_name").orderBy(col("transaction_amount").desc())
ranked_transactions = transactions_with_area.withColumn(
"row_number", row_number().over(window_spec)
)
top_5_transactions = ranked_transactions.filter(col("row_number") <= 5)
top_5_transactions.show()
Here is the output:
This approach ensures we get the top 5 transactions by amount for each area_name.
DENSE_RANK() Function
The DENSE_RANK() function returns a rank for each row in a partition, similar to RANK(), but does not leave gaps in the ranking sequence when there are ties among rows.
For example, how do we sequentially rank transactions in each area so that for tied transactions, we do not have rank gaps?
To do this, we join the transaction_records with the store to get the area_name for each transaction. Next, we partition by area_name, order by transaction_amount descending, and then use the DENSE_RANK() to assign ranks. Here’s the code:
from pyspark.sql.window import Window
from pyspark.sql.functions import col, dense_rank
transactions_with_area = transaction_records.join(
stores, transaction_records.store_id == stores.store_id
).select(
"transaction_id", "customer_id", "transaction_amount", "transaction_date", "area_name"
)
window_spec = Window.partitionBy("area_name").orderBy(col("transaction_amount").desc())
dense_ranked_transactions = transactions_with_area.withColumn(
"dense_rank", dense_rank().over(window_spec)
)
dense_ranked_transactions.show()
Here is the output:
This provides the rankings for transactions in each area_name, ensuring that tied transactions share the same rank without leaving gaps.
NTILE(n) Function
NTILE(n) is a function that distributes the sorted rows into n buckets of roughly equal size and gives each row a bucket number. It helps make quantiles or bins.
A great example is how we split transactions into four quartile buckets of sum per area.
To do this, we join the transaction_records table to the stores table on store_id to get transaction area_name. Next, we partition the data by area_name order by transaction_amount in descending order and apply the NTILE(4) function. Here’s the code:
from pyspark.sql.window import Window
from pyspark.sql.functions import col, ntile
transactions_with_area = transaction_records.join(
stores, transaction_records.store_id == stores.store_id
).select(
"transaction_id", "customer_id", "transaction_amount", "transaction_date", "area_name"
)
window_spec = Window.partitionBy("area_name").orderBy(col("transaction_amount").desc())
quartile_transactions = transactions_with_area.withColumn(
"quartile", ntile(4).over(window_spec)
)
quartile_transactions.show()
Here is the output:
This divides the transactions in each area_name into four groups, assigning each a quartile value from 1 to 4.
PySpark Analytical Functions
In this section, we will explore some PySpark analytical functions, such as the LEAD() function.
LEAD() Function
The LEAD() function retrieves the value of a column for the following row within the same window. For instance, how can we get the next transaction amount for each transaction in each area?
from pyspark.sql.window import Window
from pyspark.sql.functions import col, lead
lead_window_spec = Window.partitionBy("area_name").orderBy(col("transaction_amount").desc())
transactions_with_next = transaction_records.join(
stores, transaction_records.store_id == stores.store_id
).select(
"transaction_id", "transaction_amount", "area_name"
).withColumn(
"next_amount", lead("transaction_amount").over(lead_window_spec)
)
transactions_with_next.show()
Here is the output:
LAG() Function
The LAG() function retrieves the value of a column for the previous row within the same window. For instance, how can we get the previous transaction amount for each transaction in each area?
from pyspark.sql.functions import cume_dist
from pyspark.sql.window import Window
from pyspark.sql.functions import col,lag
lag_window_spec = Window.partitionBy("area_name").orderBy(col("transaction_amount").desc())
transactions_with_previous = transaction_records.join(
stores, transaction_records.store_id == stores.store_id
).select(
"transaction_id", "transaction_amount", "area_name"
).withColumn(
"previous_amount", lag("transaction_amount").over(lag_window_spec)
)
transactions_with_previous.show()
Here is the output:
Cumulative Distribution Function
The CUME_DIST() function calculates the cumulative distribution of a value relative to the entire window. For instance, how can we calculate the cumulative distribution of transaction amounts in each area?
from pyspark.sql.window import Window
from pyspark.sql import functions as F
# Merge transaction_records with stores to include area_name
merged_data = transaction_records.join(
stores.select("store_id", "area_name"), on="store_id", how="inner"
)
window_spec = Window.partitionBy("area_name").orderBy("transaction_amount")
result = merged_data.withColumn("cume_dist", F.cume_dist().over(window_spec))
result = result.select("transaction_id", "transaction_amount", "area_name", "cume_dist")
result.show(20)
Here is the output:
PERCENT_RANK() Function
The PERCENT_RANK() function computes the relative rank of a row within its partition as a percentage. For instance, how can we calculate the percent rank of transactions by their amounts within each area?
from pyspark.sql.functions import cume_dist
from pyspark.sql.window import Window
from pyspark.sql.functions import col
from pyspark.sql.functions import percent_rank
percent_rank_window_spec = Window.partitionBy("area_name").orderBy(col("transaction_amount").desc())
transactions_with_percent_rank = transaction_records.join(
stores, transaction_records.store_id == stores.store_id
).select(
"transaction_id", "transaction_amount", "area_name"
).withColumn(
"percent_rank", percent_rank().over(percent_rank_window_spec)
)
transactions_with_percent_rank.show()
Here is the output:
PySpark Aggregate Functions
Aggregate functions such as SUM(), AVG(), MIN(), and MAX() allow you to compute summary statistics within a partition. For this section, let’s look at shorter questions, related code, and outputs because they are straightforward.
SUM(): Total Transaction Amount Per Area
How can we calculate the total transaction amount for each area?
from pyspark.sql.functions import sum
total_transactions = transaction_records.join(
stores, transaction_records.store_id == stores.store_id
).groupBy("area_name").agg(
sum("transaction_amount").alias("total_transaction_amount")
)
total_transactions.show()
Here is the output.
AVG() – Average Transaction Amount Per Area
How can we calculate the average transaction amount for each area?
from pyspark.sql.functions import avg
average_transactions = transaction_records.join(
stores, transaction_records.store_id == stores.store_id
).groupBy("area_name").agg(
avg("transaction_amount").alias("average_transaction_amount")
)
average_transactions.show()
Here is the output:
MIN(): Minimum Transaction Amount Per Area
How can we find the smallest transaction amount in each area?
from pyspark.sql.functions import min
min_transactions = transaction_records.join(
stores, transaction_records.store_id == stores.store_id
).groupBy("area_name").agg(
min("transaction_amount").alias("min_transaction_amount")
)
min_transactions.show()
Here is the output:
MAX(): Maximum Transaction Amount Per Area
How can we find the largest transaction amount in each area?
from pyspark.sql.functions import max
max_transactions = transaction_records.join(
stores, transaction_records.store_id == stores.store_id
).groupBy("area_name").agg(
max("transaction_amount").alias("max_transaction_amount")
)
max_transactions.show()
Here is the output:
COUNT(): Number of Transactions Per Area
How can we count the number of transactions for each area?
from pyspark.sql.functions import count
transaction_count = transaction_records.join(
stores, transaction_records.store_id == stores.store_id
).groupBy("area_name").agg(
count("transaction_id").alias("transaction_count")
)
transaction_count.show()
Here is the output:
High-Density Areas
In this question, visa asks us to identify the top 3 areas with the highest customer density.
Identify the top 3 areas with the highest customer density. Customer density = (total number of unique customers in the area / area size). Your output should include the area name and its calculated customer density, and ties will be ranked the same.
Link to this question: https://platform.stratascratch.com/coding/10544-high-density-areas
In the first step, we join the transaction_records dataset along with the stores dataset. Here, we enrich the transaction data with area information (area_name and area_size). After that, we apply the simple countDistinct function to find out the total number of unique customers for each area, as explained in the analysis and aggregate functions section.
Then, we compute customer density based on unique customers/area size. This adds a calculated column, which is an example of column transformations using PySpark.
Finally, we use the dense_rank() window function to rank the areas by customer density in descending order. By grouping all areas with the same customer density, we can use the rank to take the top 3 areas and the tie as needed. Here’s the PySpark code:
from pyspark.sql import functions as F
from pyspark.sql import SparkSession
from pyspark.sql.window import Window
merged_data = transaction_records.join(stores.select('store_id', 'area_name', 'area_size'), on='store_id', how='inner')
unique_customers_per_area = merged_data.groupBy('area_name').agg(F.countDistinct('customer_id').alias('unique_customers'))
customer_density_data = unique_customers_per_area.join(stores.select('area_name', 'area_size'), on='area_name', how='inner')
customer_density_data = customer_density_data.withColumn('customer_density', F.col('unique_customers') / F.col('area_size'))
customer_density_data = customer_density_data.withColumn('rank', F.dense_rank().over(Window.orderBy(F.desc('customer_density'))))
top_3_areas = customer_density_data.filter(customer_density_data['rank'] <= 3)
final_result = top_3_areas.select('area_name', 'customer_density')
final_result_pandas = final_result.toPandas()
Here is the output:
Conclusion
In this article, we have explored different types of PySpark window functions and their explanations using real-life datasets. In the end, we solve the real interview questions.
While learning new concepts, using real-life interview questions will give you a heads-up when you land the job you want because it will become an experience.
