Implement a PySpark function to calculate the average sales per category

Instruction: Write a PySpark code snippet to calculate the average sales per category using a given DataFrame containing 'category' and 'sales' columns. Ensure your solution handles any potential data issues.

Context: This question assesses the candidate's ability to perform basic data aggregation tasks in PySpark, including grouping data and calculating averages. It also tests the candidate's awareness of data quality issues and how they might address them in their code.

Official Answer

Certainly! If I were to implement a PySpark function to calculate the average sales per category, the first step is clarifying that our DataFrame has two crucial columns: 'category' and 'sales'. My primary objective would be to ensure that the function not only performs the necessary aggregation but also gracefully handles any potential data quality issues such as missing values or incorrect data types.

To begin, let's assume our DataFrame is named salesDF. The approach I would take involves the following steps:

  1. Data Cleaning: Ensure that the 'sales' column contains numeric data and handle any null or missing values in both 'category' and 'sales' columns. For simplicity, we might decide to filter out records with null values in either column, although in a real scenario, the approach might vary based on the business context and the nature of the data.

  2. Aggregation: Group the data by the 'category' column and calculate the average sales for each category. This is a straightforward aggregation operation in PySpark, utilizing the groupBy and avg functions.

Here's how the PySpark code snippet might look:

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, avg

# Assuming SparkSession is already initialized as spark
# Initialize SparkSession (if not already initialized)
spark = SparkSession.builder.appName("AverageSalesPerCategory").getOrCreate()

# Sample DataFrame initialization (for context)
# salesDF = spark.read.csv("path_to_sales_data.csv", inferSchema=True, header=True)

# Step 1: Data Cleaning
# Filter out records with null values in 'category' or 'sales'
cleanedDF = salesDF.filter(col("category").isNotNull() & col("sales").isNotNull())

# Ensure 'sales' column is of numeric type (assuming it is already, if not, a type conversion would be necessary)

# Step 2: Aggregation
# Group by 'category' and calculate average sales
avgSalesPerCategory = cleanedDF.groupBy("category").agg(avg("sales").alias("average_sales"))

# Display the result
avgSalesPerCategory.show()

In addressing potential data issues, I first applied a filter to remove any rows with null values in 'category' or 'sales'. This is a basic yet effective way to ensure the quality of our data before performing the aggregation. Depending on the specific requirements and data characteristics, one might also consider imputing missing values or correcting data types as part of the data cleaning process.

The aggregation itself is performed using groupBy and agg functions, where avg("sales") computes the average sales, and alias("average_sales") renames the resulting column for better readability.

It's important to note that this code is designed to be adaptable. Candidates can customize the data cleaning and aggregation steps based on the specific nuances of their data and the requirements of their project. The key is to maintain a clear understanding of the data you're working with and apply PySpark's powerful data processing capabilities to derive meaningful insights efficiently.

This approach not only demonstrates the technical ability to perform data aggregation using PySpark but also shows an awareness of potential data quality issues and how to address them, which is crucial in any data engineering or data science role.

Related Questions