Lab Type

Performance Engineering

Difficulty

Expert Hands-On

Reward

800 XP · ⚡ Spark Performance Engineer

Core Stack

PySpark · Spark SQL · Spark UI

Target Goal

48 min → <3 min (16x Speedup)

⚡ Hands-On Optimization Lab

DEV-013

Optimize a Slow Spark Pipeline

A critical financial aggregation pipeline takes 48 minutes to run on a 4-node Databricks cluster, burning excessive cloud DBU credits and breaching daily SLAs. Your mission is to analyze the Spark execution plan, eliminate shuffle bottlenecks, resolve extreme data skew, replace slow Python UDFs, and achieve a 10x+ speedup.

Benchmark Baseline

Analyze Physical Plan

Eliminate Bottlenecks

Salting & Broadcast

Compare & Report

Mission & Operational Context

Understand cloud compute costs, SLA penalties, and the mechanics of Spark cluster stragglers.

The Cloud Cost Escalation

The Cloud FinOps dashboard flagged this single sales aggregation job as responsible for $4,200/month in compute spend. When looking at the Spark UI Ganglia metrics, 3 executor nodes sit at 0% CPU utilization for over 30 minutes, while 1 solitary executor core pegged at 100% CPU struggles to process a massive skewed partition.

Furthermore, the script uses a Python UDF that serializes millions of rows through Py4J, and shuffles 50 rows of store reference data into 2,000 tiny partition files. Your job as a Lead Spark Engineer is to profile, refactor, and validate this pipeline so it executes in under 3 minutes with identical results.

Executor Stragglers

When 1 partition has 17M rows and others have 50K rows, your cluster is only as fast as the slowest core.

Py4J IPC Serialization Tax

Native Spark SQL expressions execute in C++/Java bytecode via Catalyst; Python UDFs suffer massive IPC copy overhead.

Broadcast Join Power

Broadcasting tables under 100MB eliminates SortMergeJoin and avoids shuffling large datasets across the network.

Adaptive Query Execution (AQE)

Enabling AQE allows Spark to coalesce shuffle partitions and automatically handle skewed joins at runtime.

Starter Materials & Execution Plan

Inspect the slow Python workload, physical execution plan, and the partition skew profile.

📄 slow_workload.py

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, udf, when, sum, count
from pyspark.sql.types import StringType
import time

spark = SparkSession.builder \
    .appName("UnoptimizedSalesAggregation") \
    .config("spark.sql.shuffle.partitions", "2000") \
    .getOrCreate()

# -----------------------------------------------------------------
# INEFFICIENCY LAB BENCHMARK
# This job currently runs in 48 minutes on a 4-node cluster.
# Your goal is to optimize it to under 3 minutes while maintaining
# 100% identical analytical output.
# -----------------------------------------------------------------

start_time = time.time()

# 1. Inefficient Python UDF (Serialization bottleneck)
# [BOTTLENECK 1]: Python UDF forces JVM -> Python worker IPC serialization!
def categorize_tier(spend):
    if spend is None:
        return "UNKNOWN"
    spend = float(spend)
    if spend > 500:
        return "PLATINUM"
    elif spend > 200:
        return "GOLD"
    else:
        return "STANDARD"

tier_udf = udf(categorize_tier, StringType())

# 2. Large Transactions DataFrame (20M rows, skewed on STORE_ONLINE_01)
txns_df = spark.read.parquet("/data/sales_transactions_20m.parquet")

# 3. Small Store Metadata Reference Table (only 50 rows!)
# [BOTTLENECK 2]: Missing broadcast hint causes expensive SortMergeJoin and Shuffle Exchange!
stores_df = spark.read.parquet("/data/stores_reference_50_rows.parquet")

# Join large transactions with tiny stores table
joined_df = txns_df.join(stores_df, on="store_id", how="inner")

# 4. Apply slow Python UDF
enriched_df = joined_df.withColumn("loyalty_tier", tier_udf(col("net_amount")))

# 5. Multiple eager actions without caching
# [BOTTLENECK 3]: Repeated scans recompute the entire DAG from scratch!
print(f"Total Rows: {enriched_df.count()}") # Action 1: Recomputes DAG
print(f"Distinct Stores: {enriched_df.select('store_id').distinct().count()}") # Action 2: Recomputes DAG

# 6. Aggregation with severe key skew
# [BOTTLENECK 4]: 85% of volume is on 'STORE_ONLINE_01' — one single executor straggles!
summary_df = enriched_df.groupBy("store_id", "region", "loyalty_tier").agg(
    sum("net_amount").alias("total_sales"),
    count("txn_id").alias("order_count")
)

summary_df.write.format("delta").mode("overwrite").save("/lakehouse/gold/store_loyalty_summary")

print(f"Total Runtime: {time.time() - start_time:.2f} seconds")

⚡ Engineering Clue: Contains 5 major bottlenecks: Python UDF serialization, missing broadcast join on tiny 50-row table, repeated uncached actions, excessive 2000 shuffle partitions, and single-key data skew.

Hands-On Optimization Workflow

Execute these systematic optimizations to dismantle the bottlenecks one by one.

Replace Python UDF with Native Catalyst Expressions

Eliminate JVM-to-Python serialization by converting tier_udf into native when() / otherwise().

▸Replace Python UDF with: when(col('net_amount') > 500, 'PLATINUM').when(col('net_amount') > 200, 'GOLD').otherwise('STANDARD').

▸Native Spark expressions run directly in WholeStageCodegen on CPU registers with zero copy overhead.

Enable BroadcastHashJoin on Stores Table

Use broadcast(stores_df) to eliminate SortMergeJoin and shuffle exchanges.

▸stores_df has only 50 rows (<10KB). Broadcasting it replicates it to all executors in memory.

▸This completely removes the Exchange hashpartitioning step from the physical plan.

Resolve Severe Key Skew with Salting or AQE

Mitigate the 17M row skew on STORE_ONLINE_01.

▸Option A: Enable Adaptive Query Execution with spark.sql.adaptive.skewJoin.enabled = true.

▸Option B: Salt the skewed key by adding a random salt suffix (0 to 7) to distribute across multiple cores during join/aggregation.

Tune Shuffle Partitions & Remove Eager Actions

Eliminate wasteful actions and set appropriate partition sizes.

▸Remove redundant intermediate count() calls that force un-cached re-computation of the DAG.

▸Tune spark.sql.shuffle.partitions to match core count (e.g. 16 to 32) instead of the excessive 2000.

▸If intermediate results are needed multiple times, persist with .persist(StorageLevel.MEMORY_AND_DISK).

Optimize Delta Lake File Layout

Enable Liquid Clustering or Z-ORDER on the final Gold Delta table.

▸Write output partitioned or clustered by store_id and region to accelerate downstream queries.

▸Run OPTIMIZE / VACUUM checks to ensure right-sized Delta files (128MB to 1GB).

Target Performance Benchmarks

Measure before and after metrics to prove tangible engineering improvement.

Target Performance Metrics

Baseline Runtime

48.2 Minutes

2,892 Seconds

Target Optimized Runtime

< 3.0 Minutes

Target < 180 Seconds

Speedup Factor

> 16x Faster

Dramatic latency reduction

Result Reconciliation

100% Identical

Zero metric deviation

Physical Plan Transformation Comparison

Component	Unoptimized Baseline	Optimized Implementation	Engineering Rationale
Join Strategy	SortMergeJoin (Shuffled)	BroadcastHashJoin	Eliminates cross-network shuffle of 20M rows
Tier Classification	BatchEvalPython (UDF)	Native when().otherwise()	Runs in WholeStageCodegen JVM bytecode
Shuffle Partitions	2,000 tiny partitions	32 partitions / AQE auto	Reduces task scheduling latency by 95%
Data Skew Handling	1 core straggling (17M rows)	Key Salting / AQE Skew Join	Evenly saturates all 4 cluster nodes

Validation Checks & Correctness Verification

Optimization is meaningless if the data is corrupted. Prove 100% numerical equality.

Python Assertion Suite — Verify Output Correctness

# 1. Assert row counts match baseline exactly
baseline_count = baseline_df.count()
optimized_count = optimized_df.count()
assert baseline_count == optimized_count, f"Row count mismatch: {baseline_count} vs {optimized_count}"

# 2. Assert zero difference across all aggregated values
diff_df = baseline_df.subtract(optimized_df)
assert diff_df.count() == 0, "Discrepancies found between baseline and optimized tables!"

print("✓ Correctness verified! Optimized pipeline output is 100% identical to baseline.")

Evidence To Submit & Performance Report

Compile your benchmarking data and submit your final engineering report.

To receive grading and claim your badge, compile the following proof artifacts:

☑

1. Optimized Python Script (optimized_workload.py)

Clean script with broadcast join, native when/otherwise expressions, and tuned partition settings.

☑

2. Spark UI Physical Plan Screenshot

Screenshot from Spark UI DAG visualization proving BroadcastHashJoin replaced SortMergeJoin.

☑

3. Runtime Benchmark Comparison

Timing output showing baseline runtime (~48 min) vs. optimized runtime (< 3 min).

☑

4. Correctness Assertion Proof

Log output of subtract() query confirming zero variance between baseline and optimized datasets.

☑

5. Performance Engineering Summary Report (2 paragraphs)

Explanation of the biggest bottleneck, why broadcasting helped, and how key skew was mitigated.

Common Mistakes To Avoid

Review these frequent pitfalls in Spark optimization.

⚠️ Broadcasting Large DataFrames (> 100MB)

Broadcasting tables that exceed driver memory causes OutOfMemory (OOM) driver crashes. Only broadcast small lookup tables.

⚠️ Blindly Over-Partitioning Small Datasets

Setting spark.sql.shuffle.partitions = 2000 on small datasets generates thousands of tiny tasks with more scheduling overhead than actual computation.

⚠️ Assuming Caching is Always Faster

Calling .cache() on DataFrames used only once adds serialization overhead and wastes executor RAM without any performance benefit.

Stretch / Bonus Objectives

Explore cutting-edge Spark optimization features.

⭐ Dynamic Partition Pruning (DPP)

Structure queries to trigger Dynamic Partition Pruning when filtering partitioned tables via dimension joins.

⭐ Delta Liquid Clustering

Replace traditional date/store Hive partitioning with Delta Liquid Clustering and benchmark file compaction.

Optimization Certification

Complete DEV-013 to Earn 800 XP & Spark Performance Engineer Badge

Profile the slow workload, apply targeted Spark optimizations, prove 10x+ speedup with 100% numerical equality, and submit your report.

⚡

800 XP

Spark Performance Engineer