Interview Focus: PerformanceDifficulty: Hard

Optimize a Slow Spark Job

Estimated duration: 20 mins

Lab Overview & Scenario

Analyze logical query plans, use caching to prevent redundant execution, and partition data.

Business Scenario

A production ETL job is taking too long to run. You must analyze the execution plan, cache intermediate results that are reused multiple times, and partition the output by country to optimize downstream queries.

1. Local Raw Mock Dataset

The processing pipeline will ingest the following local mock file structures (e.g. simulated as `dataset.csv`):

id	user_id	action	country
1	1001	click	USA
2	1002	view	CAN
3	1001	purchase	USA
4	1003	click	GER
5	1002	click	CAN

View Raw CSV Dataset

id,user_id,action,country
1,1001,click,USA
2,1002,view,CAN
3,1001,purchase,USA
4,1003,click,GER
5,1002,click,CAN

2. Your Tasks

Read the raw user activity CSV logs.

Filter user actions related to purchases.

Cache the intermediate filtered DataFrame to prevent recalculation during multiple downstream tasks.

Save the output partitioned by the country column.

Generate and view the physical execution plan using explain().

3. Pipeline Skeleton Script

Copy this starter template to write the data transformation logic:

python

# starter.py - Optimize a Slow Spark Job
from pyspark.sql import SparkSession
from pyspark.sql.functions import col

def optimize_pipeline():
    spark = SparkSession.builder.appName("OptimizationLab").getOrCreate()
    # TODO: Read dataset.csv, filter purchases, cache, partition by country, explain plan
    pass

if __name__ == "__main__":
    optimize_pipeline()

4. Interactive Solution & Execution Output

Reveal Solution Pipeline Code

Implement the correct transformations using PySpark to achieve the aggregated values metrics output.

Pipeline Flow

dataset.csv

↓

Read CSV

↓

Filter Clicks

↓

Cache clicks_df

↓

Action 1: count()

↓

Action 2: write.partitionBy()

↓

Explain Plan

Step-by-Step Implementation Guide

Step 1

Load the activities dataset.

Step 2

Filter action logs to keep only click events.

Step 3

Call .cache() on the intermediate DataFrame. Caching keeps the computed records in executor memory, avoiding rebuilding the execution DAG when counting records and writing to Parquet.

Step 4

Save the output to Parquet partitioned by country: clicks_df.write.mode("overwrite").partitionBy("country").parquet("optimized_output.parquet").

Step 5

Invoke .explain() to display the execution plan in stdout.

Common Mistakes

Caching every single DataFrame. This wastes memory and clutters executor storage.
Forgetting to call an action after calling cache(). Caching is lazy, so the data isn't cached until the first action runs.

Interviewer's Expectations

Why this is asked:

Tests knowledge of Spark's execution graph, lazy evaluation, caching, and partitioning strategies.

What they look for:

Strategic placement of caching and write partitioning.

Common mistakes:

Caching DataFrames that are only used once, causing cache management overhead.

An excellent answer includes:

Explain that caching is lazy and discuss partition pruning benefits of saving data with partitionBy.

✔ What You Learned

✔Leverage caching on reusable DataFrames
✔Write partitioned datasets using partitionBy()
✔Inspect execution DAGs using explain()

Continue Learning

Previous LabCustomer Purchase Trends Related Theory LessonPersistence Next LabHandle Data Skew