Image by Author | Canva

Filtering data is one of the basics of data-related coding tasks because you need to filter the data for any situation.

From concepts to running a real-life interview problem from Spotify, we will cover all the steps in the article step by step until we achieve the goal! In this way, you will evolve and practice at the same time. So, let’s get started!

Setting Up PySpark

In this section, let’s see how to install PySpark and create a PySpark session.

Installing and Importing PySpark

Now let’s first install the pyspark library.

pip install pyspark

Next, let’s load the spark session from this library.

from pyspark.sql import SparkSession

Creating a Spark Session

Now, you can create a spark session like this.

spark = SparkSession.builder \
    .appName("Data Filtering with PySpark") \
    .getOrCreate()

However, if you have a Stratascratch account, you can skip the steps above because you can run pyspark code on our platform.

Loading the Dataset

You can load the dataset like this.

file_path = "path_to_your_dataset"  
df = spark.read.csv(file_path, header=True, inferSchema=True)

Viewing the Dataset

To confirm that the dataset loaded successfully, display the first few rows:

df.show()

Here is the code where you can load the dataset in strata.

from pyspark.sql import functions as F

penetration_analysis.show()

Here is the output.

Basic Data Filtering

Let’s start with basic data filtering. We will use a Spotify dataset in this section and the upcoming sections.

Spotify Penetration Analysis
Market penetration is an important metric for understanding Spotify’s performance and growth potential in different regions.
You are part of the analytics team at Spotify and are tasked with calculating the active user penetration rate in specific countries.
For this task, ‘active_users’ are defined based on the following criterias:
last_active_date: The user must have interacted with Spotify within the last 30 days.
• sessions: The user must have engaged with Spotify for at least 5 sessions.
• listening_hours: The user must have spent at least 10 hours listening on Spotify.
Based on the condition above, calculate the active ‘user_penetration_rate’ by using the following formula.
• Active User Penetration Rate = (Number of Active Spotify Users in the Country / Total users in the Country)
Total Population of the country is based on both active and non-active users.
The output should contain ‘country’ and ‘active_user_penetration_rate’ rounded to 2 decimals.
Let’s assume the current_day is 2024-01-31.

Link to this question.

Here are the columns of the dataset:

Here is a preview of the dataset:

Filtering Rows with a Single Condition

Now, let’s filter for rows where sessions exceed 5. Here is the code:

filtered_df = penetration_analysis.filter(penetration_analysis['sessions'] > 5)
filtered_df.show(5)

Here is the output.

Filtering Rows with a Multiple Condition

Let’s add one more condition. Here, let’s filter listening hours, too; here is the code:

filtered_df = penetration_analysis.filter(
    (penetration_analysis['sessions'] > 5) & 
    (penetration_analysis['listening_hours'] > 10)
)
filtered_df.show(5)

Here is the output.

Now that you understand how it works let’s use more advanced filtering.

Advanced-Data Filtering with Multiple Conditions

What if you wonder how to find genuinely engaged users with a platform? For example, those who have listened for long hours and logged frequent sessions within a recent timeframe.

The penetration_analysis dataset includes listening hours, sessions, and last_active_date, which are perfect columns for filtering for deeper insights.

Combining Time-Based and Numeric Conditions

Imagine you want to identify users who:

• Logged into the platform in the last 30 days.
• Spent over 10 hours listening.
• Had more than five sessions.

Here’s how to filter for such users:

filtered_df = penetration_analysis.filter(
    (penetration_analysis['sessions'] > 5) & 
    (penetration_analysis['listening_hours'] > 10)
)
filtered_df.show(5)

Here is the output.

Aggregations After Filtering

What if you wonder how user activity varies across countries? Aggregations help you summarize filtered data to uncover patterns and trends. With PySpark, you can group data by specific columns and apply functions like sum, average, or count for deeper analysis.

Grouping and Summing Data

Let’s calculate the total listening hours and sessions for each country. Here’s how you can do it:

from pyspark.sql import functions as F

aggregated_df = penetration_analysis.groupBy("country").agg(
    F.sum("listening_hours").alias("total_listening_hours"),
    F.sum("sessions").alias("total_sessions")
)
aggregated_df.show(5)

Here is the output:

Adding Average Metrics

To add more insights, calculate the average listening hours and sessions per user in each country:

from pyspark.sql import functions as F

aggregated_avg_df = penetration_analysis.groupBy("country").agg(
    F.avg("listening_hours").alias("average_listening_hours"),
    F.avg("sessions").alias("average_sessions")
)
aggregated_avg_df.show(5)

Here is the output.

Joining Aggregations

To combine total and average metrics, you can join these DataFrames:

final_aggregated_df = aggregated_df.join(aggregated_avg_df, on="country", how="inner") final_aggregated_df.show(5)

Here is the output:

Spotify Penetration Analysis

Good. Now we are ready, so let’s remember the question we used in its dataset from the beginning, which was asked by Spotify.

Setting Up the Problem

In this question, Spotify wants us to calculate the user_penetration_rate by using this formula:

The penetration rate formula is this: Active User Penetration Rate = (Number of Active Spotify Users in the Country / Total users in the Country)

But before this, we should filter active users with these conditions:

1. Users must have interacted with Spotify within the last 30 days.
2. They must have at least five sessions.
3. They must have spent over 10 hours listening.

To do that, we have to first:

• Filter Active Users
• Group Data and Calculate Metrics
• Calculate the Penetration Rate
• Convert to Pandas DataFrame

Filtering Active Users

To identify active users, we create a new column, is_active, which evaluates whether a user meets all three conditions:

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, count, sum as sum_, round as round_
from pyspark.sql.types import DateType
from datetime import datetime, timedelta

# Fixed reference date
fixed_current_date = datetime(2024, 1, 31)
within_last_30_days = (fixed_current_date - timedelta(days=30)).date()

# Add "is_active" column to indicate active users
penetration_analysis = penetration_analysis.withColumn(
    "is_active",
    (col("last_active_date") >= within_last_30_days) &
    (col("sessions") >= 5) &
    (col("listening_hours") >= 10)
).show()

Here is the output.

Grouping Data and Calculating Metrics

Next, we group the dataset by country and calculate two metrics:

1. Total users: Count of all rows per country.
2. Active users: Sum of is_active (casting True to 1 and False to 0).

grouped = penetration_analysis.groupBy("country").agg(
    count("is_active").alias("total_users"),
    sum_(col("is_active").cast("int")).alias("active_users")
).show()

Here is the output.

Calculating the Penetration Rate

Finally, calculate the penetration rate by dividing active_users by total_users and round the result to two decimals:

result = grouped.withColumn(
    "penetration",
    round_(col("active_users") / col("total_users"), 2)
).select("country", "penetration")

Here is the output.

Converting to Pandas for Output

In the end, we will convert the result to the output.

result = result.toPandas()

Here is the entire code:

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, count, sum as sum_, round as round_
from pyspark.sql.types import DateType
from datetime import datetime, timedelta

# Setting a fixed current date for the analysis
fixed_current_date = datetime(2024, 1, 31)
within_last_30_days = (fixed_current_date - timedelta(days=30)).date()

# Convert 'last_active_date' to DateType if necessary
penetration_analysis = penetration_analysis.withColumn("last_active_date", col("last_active_date").cast(DateType()))

# Filter the DataFrame for active users
penetration_analysis = penetration_analysis.withColumn(
    "is_active",
    (col("last_active_date") >= within_last_30_days) &
    (col("sessions") >= 5) &
    (col("listening_hours") >= 10)
)

# Grouping and calculating the penetration
grouped = penetration_analysis.groupBy("country").agg(
    count("is_active").alias("count"),
    sum_(col("is_active").cast("int")).alias("sum")
)

result = grouped.withColumn("penetration", round_(col("sum") / col("count"), 2)).select("country", "penetration")

result = result.toPandas()

Here is the output:

You can also see a walkthrough of each question like this:

Conclusion

In this article, we have explored different data filtering use cases, starting from basic and going through advanced filtering options. In the end, we solved the interview question from Spotify. So, by going through it, you have learned data filtering with Pyspark and prepared yourself for the interviews. With this practical example, you’re now better equipped to apply these skills in your data science journey and future interviews.