GroupBy and aggregations in Apache Spark - Practice Problems & Coding Challenges

Challenge - 5 Problems

❓ Predict Output

intermediate

Output of GroupBy with sum aggregation

What is the output of the following Apache Spark code snippet?

Apache Spark

from pyspark.sql import SparkSession
from pyspark.sql.functions import sum

spark = SparkSession.builder.appName('test').getOrCreate()
data = [(1, 'A', 10), (2, 'B', 20), (3, 'A', 30), (4, 'B', 40)]
columns = ['id', 'category', 'value']
df = spark.createDataFrame(data, columns)
result = df.groupBy('category').agg(sum('value').alias('total_value')).orderBy('category')
result.show()

A [Row(category='A', total_value=10), Row(category='B', total_value=20)]

B [Row(category='A', total_value=30), Row(category='B', total_value=40)]

C [Row(category='A', total_value=40), Row(category='B', total_value=60)]

D [Row(category='A', total_value=70), Row(category='B', total_value=60)]

💡 Hint

Sum the 'value' column grouped by 'category'.

The sum aggregation adds values for each category: for 'A' it's 10 + 30 = 40, for 'B' it's 20 + 40 = 60.

❓ data_output

intermediate

Count distinct values per group

Given the DataFrame below, what is the output of counting distinct 'value' per 'category'?

Apache Spark

from pyspark.sql import SparkSession
from pyspark.sql.functions import countDistinct

spark = SparkSession.builder.appName('test').getOrCreate()
data = [(1, 'A', 10), (2, 'B', 20), (3, 'A', 10), (4, 'B', 40), (5, 'A', 30)]
columns = ['id', 'category', 'value']
df = spark.createDataFrame(data, columns)
result = df.groupBy('category').agg(countDistinct('value').alias('distinct_values')).orderBy('category')
result.show()

A [Row(category='A', distinct_values=2), Row(category='B', distinct_values=3)]

B [Row(category='A', distinct_values=2), Row(category='B', distinct_values=2)]

C [Row(category='A', distinct_values=3), Row(category='B', distinct_values=3)]

D [Row(category='A', distinct_values=3), Row(category='B', distinct_values=2)]

💡 Hint

Count unique 'value' entries per 'category'.

Category 'A' has values 10, 10, 30 → 2 distinct values; 'B' has 20, 40 → 2 distinct values.

❓ visualization

advanced

Visualizing average values per group

Which option correctly describes the bar chart produced by the following code?

Apache Spark

import matplotlib.pyplot as plt
from pyspark.sql import SparkSession
from pyspark.sql.functions import avg

spark = SparkSession.builder.appName('test').getOrCreate()
data = [(1, 'X', 5), (2, 'Y', 15), (3, 'X', 10), (4, 'Y', 20), (5, 'Z', 25)]
columns = ['id', 'group', 'score']
df = spark.createDataFrame(data, columns)
avg_df = df.groupBy('group').agg(avg('score').alias('avg_score')).orderBy('group')
pandas_df = avg_df.toPandas()
plt.bar(pandas_df['group'], pandas_df['avg_score'])
plt.xlabel('Group')
plt.ylabel('Average Score')
plt.title('Average Score per Group')
plt.show()

A Bar chart with groups X, Y, Z on x-axis and average scores 7.5, 17.5, 25 respectively on y-axis.

B Bar chart with groups X, Y, Z on x-axis and total scores 15, 35, 25 respectively on y-axis.

C Line chart with groups X, Y, Z on x-axis and average scores 7.5, 17.5, 25 respectively on y-axis.

D Bar chart with groups X, Y, Z on x-axis and median scores 7.5, 17.5, 25 respectively on y-axis.

💡 Hint

The code uses avg aggregation and bar chart.

The average scores are (5+10)/2=7.5 for X, (15+20)/2=17.5 for Y, and 25 for Z. The chart is a bar chart.

🔧 Debug

advanced

Identify the error in aggregation code

What error will the following Apache Spark code raise?

Apache Spark

from pyspark.sql import SparkSession
from pyspark.sql.functions import mean

spark = SparkSession.builder.appName('test').getOrCreate()
data = [(1, 'A', 10), (2, 'B', 20)]
columns = ['id', 'category', 'value']
df = spark.createDataFrame(data, columns)
result = df.groupBy('category').agg(mean('values').alias('avg_value'))
result.show()

A AnalysisException: cannot resolve '`values`' given input columns: [id, category, value]

B AttributeError: 'DataFrame' object has no attribute 'mean'

C No error, outputs average values per category

D TypeError: mean() missing 1 required positional argument

💡 Hint

Check the column name used in mean function.

The code uses 'values' instead of 'value' which does not exist in DataFrame columns, causing AnalysisException.

🚀 Application

expert

Calculate weighted average per group

You have a DataFrame with columns 'group', 'score', and 'weight'. Which code snippet correctly calculates the weighted average score per group?

Apache Spark

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

spark = SparkSession.builder.appName('test').getOrCreate()
data = [('A', 80, 0.5), ('A', 90, 0.5), ('B', 70, 0.3), ('B', 60, 0.7)]
columns = ['group', 'score', 'weight']
df = spark.createDataFrame(data, columns)

A df.groupBy('group').agg((sum('score' * 'weight') / sum('weight')).alias('weighted_avg')).show()

B df.groupBy('group').agg((sum(col('score') * col('weight')) / sum('weight')).alias('weighted_avg')).show()

C df.groupBy('group').agg((sum(col('score') * col('weight')) / sum('weight')).alias('weighted_avg')).collect()

D df.groupBy('group').agg((sum(col('score') * col('weight')) / sum(col('weight'))).alias('weighted_avg')).show()

💡 Hint

Use col() for column operations inside sum().

Option D correctly uses col() for both numerator and denominator sums and shows the result. Option D misses col() in denominator causing error. Option D uses strings incorrectly. Option D collects but does not show output.

Practice

1. What does the groupBy function do in Apache Spark?

easy

A. It splits the data into groups based on column values.

B. It sorts the data in ascending order.

C. It filters rows based on a condition.

D. It joins two DataFrames together.

Solution

1. Step 1: Understand the purpose of groupBy

   The groupBy function is used to split data into groups based on one or more column values.

2. Step 2: Compare with other functions

   Sorting, filtering, and joining are different operations and do not split data into groups.

3. Final Answer:

   It splits the data into groups based on column values. -> Option A

4. Quick Check:

   groupBy = splits data into groups [OK]

Hint: Remember: groupBy groups rows by column values [OK]

Common Mistakes:

• Confusing groupBy with filter or sort
• Thinking groupBy changes row order
• Assuming groupBy merges DataFrames

2. Which of the following is the correct syntax to calculate the average of column sales after grouping by region in Spark DataFrame df?

easy

A. df.groupBy("region").sum("sales")

B. df.groupBy("region").agg(avg("sales"))

C. df.groupBy("region").filter(avg("sales"))

D. df.groupBy("region").select(avg("sales"))

Solution

1. Step 1: Identify correct aggregation syntax

   To calculate average after grouping, use groupBy followed by agg(avg(column)).

2. Step 2: Check each option

   df.groupBy("region").agg(avg("sales")) uses agg(avg("sales")) correctly. df.groupBy("region").sum("sales") uses sum, not average. df.groupBy("region").filter(avg("sales")) uses filter incorrectly. df.groupBy("region").select(avg("sales")) uses select which is not valid after groupBy for aggregation.

3. Final Answer:

   df.groupBy("region").agg(avg("sales")) -> Option B

4. Quick Check:

   groupBy + agg(avg) = df.groupBy("region").agg(avg("sales")) [OK]

Hint: Use agg() with avg() inside groupBy for averages [OK]

Common Mistakes:

• Using sum instead of avg for average
• Using filter instead of agg
• Using select after groupBy without agg

3. Given the following Spark code, what will be the output schema and result?

from pyspark.sql import SparkSession
from pyspark.sql.functions import sum

spark = SparkSession.builder.getOrCreate()
data = [("East", 100), ("West", 200), ("East", 150)]
df = spark.createDataFrame(data, ["region", "sales"])
result = df.groupBy("region").agg(sum("sales").alias("total_sales"))
result.show()

medium

A. region | total_sales
East | 250
West | 200

B. region | sales
East | 100
West | 200
East | 150

C. region | total_sales
East | 150
West | 200

D. region | total_sales
East | 100
West | 200

Solution

1. Step 1: Understand groupBy and sum aggregation

   The code groups rows by "region" and sums the "sales" values for each group.

2. Step 2: Calculate sums for each region

   For "East": 100 + 150 = 250; For "West": 200.

3. Final Answer:

   region | total_sales East | 250 West | 200 -> Option A

4. Quick Check:

   Sum sales by region = 250 (East), 200 (West) [OK]

Hint: Sum sales per group to get total_sales [OK]

Common Mistakes:

• Not summing all values per group
• Confusing alias name with original column
• Expecting original rows without aggregation

4. Identify the error in the following Spark code snippet:

df.groupBy("category").agg(sum("amount"), avg("amount"))

medium

A. Missing alias for aggregated columns causes error.

B. agg() requires a dictionary, not multiple arguments.

C. groupBy cannot be used with agg.

D. sum and avg functions are not imported.

Solution

1. Step 1: Check function imports

   The code uses sum and avg functions but does not show imports from pyspark.sql.functions.

2. Step 2: Validate other options

   agg() accepts multiple arguments; aliasing is optional; groupBy works with agg.

3. Final Answer:

   sum and avg functions are not imported. -> Option D

4. Quick Check:

   Missing imports cause NameError [OK]

Hint: Always import aggregation functions before use [OK]

Common Mistakes:

• Assuming agg needs dict, but it accepts multiple args
• Forgetting to import sum, avg from pyspark.sql.functions
• Thinking groupBy can't be chained with agg

5. You have a Spark DataFrame sales_df with columns store, product, and revenue. You want to find the maximum revenue per store and the average revenue per product across all stores in a single query. Which code snippet correctly achieves this?

hard

A. sales_df.groupBy("store").agg(max("revenue").alias("max_revenue"), avg("revenue").alias("avg_revenue"))

B. sales_df.agg(max("revenue").alias("max_revenue"), avg("revenue").alias("avg_revenue"))

C. sales_df.groupBy("store").agg(max("revenue").alias("max_revenue")).crossJoin(sales_df.groupBy("product").agg(avg("revenue").alias("avg_revenue")))

D. sales_df.groupBy("store").agg(max("revenue").alias("max_revenue")).groupBy("product").agg(avg("revenue").alias("avg_revenue"))

Solution

1. Step 1: Understand the requirements

   We need max revenue per store and average revenue per product in one result.

2. Step 2: Analyze options

   sales_df.groupBy("store").agg(max("revenue").alias("max_revenue"), avg("revenue").alias("avg_revenue")) aggregates both by store only, so avg revenue per product is incorrect. sales_df.agg(max("revenue").alias("max_revenue"), avg("revenue").alias("avg_revenue")) aggregates globally, losing grouping. sales_df.groupBy("store").agg(max("revenue").alias("max_revenue")).groupBy("product").agg(avg("revenue").alias("avg_revenue")) tries chaining groupBy but overwrites previous aggregation. sales_df.groupBy("store").agg(max("revenue").alias("max_revenue")).crossJoin(sales_df.groupBy("product").agg(avg("revenue").alias("avg_revenue"))) computes both groupings separately and joins results, correctly combining both summaries.

3. Final Answer:

   sales_df.groupBy("store").agg(max("revenue").alias("max_revenue")).crossJoin(sales_df.groupBy("product").agg(avg("revenue").alias("avg_revenue"))) -> Option C

4. Quick Check:

   Separate groupBy + crossJoin combines both aggregations [OK]

Hint: Use separate groupBy then join to combine different aggregations [OK]

Common Mistakes:

• Trying to aggregate different groupings in one groupBy
• Overwriting aggregation by chaining groupBy
• Forgetting to join separate aggregation results