Diagnose and Fix a Slow PySpark Pipeline

Tune a PySpark batch pipeline by reading execution plans and fixing shuffles.

Introduction

30 Second Summary

Every online store runs data pipelines that crunch thousands of orders into revenue reports and customer insights. When those pipelines slow down, the entire business waits.

In this project, you will build a PySpark batch pipeline that processes one million e-commerce orders, intentionally trigger a slow join, and fix it. Along the way, you will learn to read Spark execution plans and apply broadcast joins and Adaptive Query Execution.

What You'll Build

You run your pipeline and a revenue-by-category breakdown appears in your terminal, built from one million orders joined with a product catalog using an execution plan you diagnosed and fixed yourself.

By the end of this project, you'll have:

A working PySpark batch pipeline that reads one million e-commerce orders, joins them with a product catalog, and prints a revenue-by-category breakdown to your terminal.
The ability to read a Spark execution plan and identify the shuffle operations that slow pipelines at scale.
An optimized pipeline using broadcast joins, AQE, and partitioned Parquet output ready for downstream consumption.
Secret Mission: Deploy your optimized pipeline to AWS EMR Serverless and run a Spark job in the cloud.

Are there any prerequisites?

You should be comfortable with Python and DataFrames (from pandas or a similar library). No prior Spark experience is required.

Before We Start

Before writing any code, let's lock in what you're about to build and why it matters.

Install Java and PySpark

PySpark is Python's interface to Apache Spark. Under the hood, Spark runs on the Java Virtual Machine (JVM), so Java must be installed before PySpark can execute anything.

In this step, you will install Java 17 and PySpark 4.2.0. Getting both set up now means every step after this gives you instant feedback when you run your pipeline.

In this step, get ready to:

Install Java 17 (OpenJDK) and set the JAVA_HOME environment variable.
Install PySpark 4.2.0 via pip.
Verify that PySpark can start a SparkSession locally.

Install Java 17 and set JAVA_HOME

Spark 4.x requires JDK 17 or higher. Support for Java 8 and 11 was removed, so older versions will not work.

Check whether Java is already installed by running this command:

java -version

✔️ I see version 17 or higher

Java 17 is ready. Skip ahead to setting JAVA_HOME below.

ⓧ I see an older version

Spark 4.x dropped support for Java 8 and 11. Install Java 17 alongside your current version.

macOS

Install OpenJDK 17 via Homebrew by running:

brew install openjdk@17

Create a symlink so macOS can find the JDK. Enter your login password when prompted:

sudo ln -sfn $(brew --prefix)/opt/openjdk@17/libexec/openjdk.jdk /Library/Java/JavaVirtualMachines/openjdk-17.jdk

Linux

Install OpenJDK 17 by running:

sudo apt install openjdk-17-jdk

Enter your password when prompted. The install may take a minute.

Windows

Download the JDK 17 installer from adoptium.net/temurin/releases.
Select JDK 17 (LTS) and Windows as the operating system.
Run the installer and select the Set JAVA_HOME variable checkbox when prompted.

ⓧ Command not found

Java is not installed yet. Install OpenJDK 17 now.

macOS

Install OpenJDK 17 via Homebrew by running:

brew install openjdk@17

Create a symlink so macOS can find the JDK. Enter your login password when prompted:

sudo ln -sfn $(brew --prefix)/opt/openjdk@17/libexec/openjdk.jdk /Library/Java/JavaVirtualMachines/openjdk-17.jdk

Linux

Install OpenJDK 17 by running:

sudo apt install openjdk-17-jdk

Enter your password when prompted. The install may take a minute.

Windows

Download the JDK 17 installer from adoptium.net/temurin/releases.
Select JDK 17 (LTS) and Windows as the operating system.
Run the installer and select the Set JAVA_HOME variable checkbox when prompted.

What is JAVA_HOME?

JAVA_HOME is an environment variable that tells tools like PySpark where to find Java on your machine. Without it, Spark cannot locate the JVM and will fail to start.

Now set JAVA_HOME so PySpark can find your Java installation.

macOS

Add JAVA_HOME to your shell configuration by running these commands:

echo 'export JAVA_HOME=$(/usr/libexec/java_home -v 17)' >> ~/.zshrc
source ~/.zshrc

Linux

Add JAVA_HOME to your shell configuration by running these commands:

echo 'export JAVA_HOME=$(dirname $(dirname $(readlink -f $(which java))))' >> ~/.bashrc
source ~/.bashrc

Windows

The Adoptium installer sets JAVA_HOME automatically if you selected that option. Close and reopen your terminal for the change to take effect.

Confirm JAVA_HOME is set by running:

echo $env:JAVA_HOME

You should see a path containing jdk-17. If the output is empty, re-run the Adoptium installer and select the Set JAVA_HOME variable option.

Confirm Java is working by running:

java -version

You should see output starting with openjdk version "17. followed by a patch number. The exact patch number does not matter, as long as the major version is 17.

Still seeing an older Java version?

Make sure you saved your shell config file (~/.zshrc on macOS, ~/.bashrc on Linux) and ran source on it.
On macOS, verify the symlink was created after running brew install openjdk@17.
If you have multiple Java versions installed, JAVA_HOME controls which one is active. Double-check the value with echo $JAVA_HOME.

Help me troubleshoot my Java 17 installation

Install PySpark

PySpark is the Python API for Apache Spark. Installing it via pip brings down both the Python bindings and a bundled Spark runtime.

Install PySpark 4.2.0 by running:

pip install pyspark==4.2.0

What does this command do?

The pyspark package includes both the Python API and a complete Spark runtime. You do not need a separate Spark download.

The pip install puts everything PySpark needs into your Python environment.

Your terminal shows download progress and finishes with a success message.

pip install failed?

PySpark 4.2.0 requires Python 3.10 or higher. Check your version with python --version.
If your system uses python3 and pip3, try pip3 install pyspark==4.2.0 instead.
Make sure pip is up to date by running pip install --upgrade pip.

Help me fix my PySpark installation error

Before you run the next command, what do you expect to see? Will PySpark be able to start a Spark session and print a version number?

Test that PySpark and Java work together by running:

python -c "from pyspark.sql import SparkSession; spark = SparkSession.builder.master('local[*]').getOrCreate(); print(spark.version); spark.stop()"

You should see 4.2.0 printed in your terminal, along with some Spark log messages. This confirms Java, PySpark, and the Spark runtime are all connected.

Seeing a Java-related error?

Check that JAVA_HOME is set by running echo $JAVA_HOME.
If the variable is empty or points to the wrong version, go back to the JAVA_HOME setup above.
On macOS, make sure you created the symlink after running brew install openjdk@17.

PySpark cannot find Java on my system

Your environment is ready. Java and PySpark are installed and working together. Next up, you will generate a million rows of e-commerce data and build your first aggregation.

Generate Data and Build Your First Aggregation

Your environment is set up. Java and PySpark are installed, and Spark can start a local session. Now you need data to process.

In this step, you'll create two Python scripts. The first generates a million synthetic e-commerce orders and a product catalog. The second reads that data and shows you which products earn the most revenue.

In this step, get ready to:

Create a data generator that produces 1 million orders and 1,000 products as Parquet files.
Build an aggregation pipeline that calculates revenue per product and prints the top 10.

Create the project folder

You need a dedicated folder to hold your scripts and the data they generate.

Open your terminal. Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
Type Terminal (macOS) or PowerShell (Windows) and press Enter to open it.
Move to your Desktop by running this command:

cd ~/Desktop

Create your project folder and move into it by running these commands:

mkdir spark-pipeline
cd spark-pipeline

Confirm you're in the right place by running:

pwd

You should see a path ending in spark-pipeline.

Don't see spark-pipeline?

Make sure you ran both the mkdir and cd commands. If you see your home directory instead, run cd ~/Desktop/spark-pipeline to get back on track.

Help me navigate to my project folder

Create the data generator script

Your data generator will create two datasets. A small product catalog with 1,000 rows acts as a dimension table. A large orders table with 1 million rows acts as a fact table.

You'll also bake in a data skew problem on purpose. 30% of all orders will go to just 5 products. This is realistic and will matter when you diagnose performance later.

Open your text editor. Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
Type VS Code (or the name of your text editor) and press Enter to open it.
Click File in the top menu bar.
Click Open Folder.
Select the spark-pipeline folder on your Desktop, then click Open.
Create a new file called generate_data.py and add the following code to set up the Spark session and generate the products table:

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder \
    .appName("GenerateData") \
    .master("local[*]") \
    .getOrCreate()

NUM_ORDERS = 1_000_000
NUM_PRODUCTS = 1_000

# Products: small dimension table (1000 rows)
products = spark.range(1, NUM_PRODUCTS + 1).toDF("product_id")
products = products.withColumn(
    "category",
    F.when(F.col("product_id") % 5 == 0, F.lit("Electronics"))
    .when(F.col("product_id") % 5 == 1, F.lit("Clothing"))
    .when(F.col("product_id") % 5 == 2, F.lit("Books"))
    .when(F.col("product_id") % 5 == 3, F.lit("Home"))
    .otherwise(F.lit("Sports"))
)
products = products.withColumn(
    "product_name",
    F.concat(F.lit("Product_"), F.col("product_id").cast("string"))
)

products.write.mode("overwrite").parquet("data/products")

What does this code do?

SparkSession.builder creates a local Spark session. local[*] tells Spark to use all available CPU cores on your machine.
spark.range(1, NUM_PRODUCTS + 1).toDF("product_id") creates a DataFrame with 1,000 rows numbered 1 through 1,000.
The F.when(...).otherwise(...) chain assigns each product a category based on its ID. This gives you five categories: Electronics, Clothing, Books, Home, and Sports.
products.write.mode("overwrite").parquet("data/products") saves the products DataFrame as Parquet files inside a data/products directory.

What is Parquet?

Parquet is a columnar file format designed for big data. Instead of storing data row by row (like CSV), it stores data column by column. This makes it much faster for analytical queries that only need a few columns out of many.

Add the following code below the products section to generate the orders table and finish the script:

# Orders: large fact table (1M rows) with skewed product distribution
# 30% of orders go to just 5 popular products (IDs 1-5)
# 70% spread across remaining 995 products
orders = spark.range(0, NUM_ORDERS).toDF("order_id")
orders = orders.withColumn(
    "product_id",
    F.when(F.rand() < 0.3, (F.floor(F.rand() * 5) + 1).cast("int"))
    .otherwise((F.floor(F.rand() * 995) + 6).cast("int"))
)
orders = orders.withColumn(
    "quantity",
    (F.floor(F.rand() * 10) + 1).cast("int")
)
orders = orders.withColumn(
    "unit_price",
    F.round(F.rand() * 100 + 5, 2)
)

orders.write.mode("overwrite").parquet("data/orders")

print(f"Generated {NUM_ORDERS} orders across {NUM_PRODUCTS} products")
print("Skew: 30% of orders concentrated on product IDs 1-5")
print("Data saved to data/orders and data/products")

spark.stop()

What does this code do?

F.when(F.rand() < 0.3, ...) creates the skew. 30% of orders get a product_id between 1 and 5. The remaining 70% spread across product IDs 6 through 1,000.
quantity is a random integer from 1 to 10. unit_price is a random value between 5.00 and 105.00.
spark.stop() shuts down the Spark session cleanly. Always call this at the end of a PySpark script.

Why build in data skew on purpose?

In production data, some values are always more popular than others. A few bestselling products dominate order volume. This concentration is called data skew, and it causes real performance problems in Spark because one task gets stuck processing a massive partition while the others finish quickly.

By baking skew into the data now, you'll be able to see its effects when you diagnose the pipeline later.

Save generate_data.py.
Switch back to your terminal and run the script by running this command:

python generate_data.py

Spark will print some log output as it starts up. At the end, you should see these three lines:

Generated 1000000 orders across 1000 products
Skew: 30% of orders concentrated on product IDs 1-5
Data saved to data/orders and data/products

Confirm the data directories were created by running:

ls data/

You should see orders and products listed as directories.

Script failed or no data directory?

If you see a Java-related error, check that JAVA_HOME is set correctly by running echo $JAVA_HOME. It should point to your Java 17 installation.
If you see a syntax error, make sure you copied both code chunks into generate_data.py in order. The second chunk should appear directly below the first.
Make sure you are running the command from inside the spark-pipeline folder. Run pwd to check.

Help me debug my generate_data.py script

✔️ Awesome, I've got everything!

Great. Make sure you've saved generate_data.py before moving on.

ⓧ I'd like to double check the full code

Your generate_data.py should contain these two sections in order. Scroll up and confirm each one is present:

Chunk 1: Imports, SparkSession setup, constants, products generation, and products.write.mode("overwrite").parquet("data/products").
Chunk 2: Orders generation with skewed distribution, orders.write.mode("overwrite").parquet("data/orders"), print statements, and spark.stop().

Build the aggregation pipeline

Your data is ready. Now you need a script that reads it, calculates revenue per order line, and shows you which products are earning the most.

Create a new file called pipeline.py in your spark-pipeline folder and add the following code:

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder \
    .appName("EcommercePipeline") \
    .master("local[*]") \
    .getOrCreate()

# Read data
orders = spark.read.parquet("data/orders")

# Calculate revenue per order line
orders = orders.withColumn("revenue", F.col("quantity") * F.col("unit_price"))

# Aggregate revenue by product_id
product_revenue = orders.groupBy("product_id").agg(
    F.sum("revenue").alias("total_revenue")
)

# Show top 10 products by revenue
print("\n=== Top 10 Products by Revenue ===")
product_revenue.orderBy(F.col("total_revenue").desc()).show(10)

spark.stop()

What does this code do?

spark.read.parquet("data/orders") reads the Parquet files you generated and loads them as a DataFrame.
withColumn("revenue", F.col("quantity") * F.col("unit_price")) adds a new column that calculates revenue for each order line.
groupBy("product_id").agg(F.sum("revenue")) groups all orders by product and sums the revenue for each one. This is the PySpark equivalent of a SQL GROUP BY.
orderBy(F.col("total_revenue").desc()).show(10) sorts by revenue in descending order and prints the top 10 as a formatted table.

Save pipeline.py.

Before you run this, what do you expect to see? The script reads a million orders and groups them by product. Which product IDs do you think will have the highest revenue?

Run the pipeline by running this command:

python pipeline.py

After Spark's log output, you should see a formatted table under the heading === Top 10 Products by Revenue ===. The top product IDs should be 1 through 5. Those are the skewed products that received 30% of all orders.

Not seeing the table?

If you see a file-not-found error mentioning data/orders, make sure you ran python generate_data.py first. The pipeline reads the Parquet files that script creates.
If the table appears but the top products are not IDs 1-5, that's OK. The data uses random generation, so slight variation is normal. The skewed products should still dominate the top spots.
If you see a syntax error, double check that you copied all the code into pipeline.py. Make sure the parentheses in the .agg() call are balanced.

Help me debug my pipeline.py script

✔️ Awesome, I've got everything!

Great. Make sure you've saved pipeline.py before moving on.

ⓧ I'd like to double check the full code

Here is the complete pipeline.py for comparison:

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder \
    .appName("EcommercePipeline") \
    .master("local[*]") \
    .getOrCreate()

# Read data
orders = spark.read.parquet("data/orders")

# Calculate revenue per order line
orders = orders.withColumn("revenue", F.col("quantity") * F.col("unit_price"))

# Aggregate revenue by product_id
product_revenue = orders.groupBy("product_id").agg(
    F.sum("revenue").alias("total_revenue")
)

# Show top 10 products by revenue
print("\n=== Top 10 Products by Revenue ===")
product_revenue.orderBy(F.col("total_revenue").desc()).show(10)

spark.stop()

Your pipeline works. You can generate a million orders, aggregate them, and see revenue results in your terminal. But right now, the pipeline only knows product IDs. Next up, you'll join with the products table to get category-level insights, and discover the most common Spark performance problem along the way.

Add a Join and Watch It Shuffle

In the last step, you generated a million orders and built your first aggregation. Your pipeline shows revenue by product ID, but product IDs alone don't tell a business story.

To understand which product categories drive revenue, you need to join orders with the products table. This join introduces the most common Spark performance problem: a full shuffle on both sides of the join. In this step, you'll add the join, read Spark's execution plan, and see the shuffle happen firsthand.

In this step, get ready to:

Disable AQE and auto-broadcast to expose Spark's default join behavior.
Join orders with products and diagnose the execution plan.
Update the aggregation to report revenue by category.

Disable AQE and join orders with products

Spark has a built-in optimizer called Adaptive Query Execution (AQE) that automatically picks the best join strategy at runtime. It also auto-broadcasts small tables when they fit under a size threshold.

To see what a bad join plan looks like, you need to turn both of those optimizations off. This is realistic because AQE doesn't apply to streaming queries, and stale table statistics can fool even AQE in production.

In pipeline.py, find the SparkSession builder block:

spark = SparkSession.builder \
    .appName("EcommercePipeline") \
    .master("local[*]") \
    .getOrCreate()

Replace it with this version that disables AQE and auto-broadcast:

spark = SparkSession.builder \
    .appName("EcommercePipeline") \
    .master("local[*]") \
    .config("spark.sql.adaptive.enabled", "false") \
    .config("spark.sql.autoBroadcastJoinThreshold", "-1") \
    .getOrCreate()

What do these configs do?

spark.sql.adaptive.enabled set to false turns off AQE so Spark won't secretly optimize the join for you. You want to see the unoptimized plan first.

spark.sql.autoBroadcastJoinThreshold set to -1 disables automatic broadcasting of small tables. By default, Spark broadcasts any table under 10 MB. Disabling it forces Spark to use a sort-merge join instead.

Now you need to read the products table, join it with orders, and ask Spark to show you its execution plan.

Find the line orders = spark.read.parquet("data/orders") and add this line directly below it:

products = spark.read.parquet("data/products")

Find the revenue calculation line (orders = orders.withColumn("revenue", ...)). Delete everything below it (the old aggregation and print statements) and replace it with the following code:

# Join orders with products on product_id
joined = orders.join(products, "product_id")

# Inspect the physical plan - look for SortMergeJoin and Exchange nodes
print("=== Physical Plan ===")
joined.explain(True)

What does this code do?

orders.join(products, "product_id") joins every order row with the matching product row based on product_id. After this, each row has both the order details and the product's category and name.
joined.explain(True) prints Spark's full execution plan to the terminal. This shows you exactly how Spark intends to perform the join.

Before you run this, think about it: Spark has 1 million orders and 1,000 products stored in separate Parquet files. To match each order with its product, will Spark keep the data in place, or will it need to move data around?

Save pipeline.py and run it by executing this command:

python pipeline.py

Look at the physical plan output above the Spark logs. You should see SortMergeJoin and two Exchange hashpartitioning nodes in the plan. That's the shuffle.

What are SortMergeJoin and Exchange?

SortMergeJoin means Spark is sorting both datasets by the join key (product_id) and then merging them together. It's the most general-purpose join strategy, but also the most expensive.

Each Exchange hashpartitioning node means Spark is shuffling data across partitions. Two Exchange nodes means it shuffles both sides of the join. Spark is moving the entire 1-million-row orders table AND the 1,000-row products table across the network to co-locate matching rows.

Shuffling a 1,000-row table is absurd. At production scale, this shuffle is the number one performance killer in Spark pipelines.

Don't see SortMergeJoin in the output?

If you see BroadcastHashJoin instead, your config lines didn't take effect. Double-check that spark.sql.adaptive.enabled is set to "false" (as a string, not a Python boolean) and spark.sql.autoBroadcastJoinThreshold is set to "-1".
Make sure both .config() calls appear before .getOrCreate() in the builder chain.
Still stuck? Why is my Spark plan showing BroadcastHashJoin instead of SortMergeJoin even though I disabled auto-broadcast?

Update the aggregation to revenue by category

Now that each order row includes the product's category from the join, you can aggregate revenue by category instead of by product ID.

In pipeline.py, find the joined.explain(True) line you just added. Add the following code directly below it:

# Aggregate revenue by category
category_revenue = joined.groupBy("category").agg(
    F.sum("revenue").alias("total_revenue")
)

# Show results
print("\n=== Revenue by Category ===")
category_revenue.orderBy(F.col("total_revenue").desc()).show()

spark.stop()

What does this code do?

joined.groupBy("category") groups the joined data by product category instead of individual product ID. This gives you five groups (Electronics, Clothing, Books, Home, Sports).
The orderBy(F.col("total_revenue").desc()) sorts categories from highest to lowest revenue so you can see which category leads.

Before you run this, what do you expect to see? Five category rows with revenue totals, or something else?

Save pipeline.py and run the pipeline again:

python pipeline.py

You should see two things in the output: the physical plan (still showing SortMergeJoin and two Exchange nodes), followed by a table showing revenue broken down by the five categories.

Seeing a NameError or missing column?

If you see an error about category not being found, make sure the aggregation uses joined.groupBy("category") (the joined DataFrame), not orders.groupBy("category"). The category column only exists after the join.
Check that you deleted the old aggregation code (product_revenue = orders.groupBy("product_id")...). Having both the old and new aggregation will cause duplicate output or variable conflicts.
Still stuck? Help me debug my PySpark category aggregation after joining orders with products.

✔️ Awesome, I've got everything!

Great. Double-check you've saved pipeline.py and that your terminal shows both the physical plan and the revenue-by-category table.

ⓧ I'd like to double check the full code

Your pipeline.py should contain all of the following sections in order. Scroll up and confirm each one is present:

Imports (SparkSession and functions as F).
SparkSession builder with spark.sql.adaptive.enabled set to "false" and spark.sql.autoBroadcastJoinThreshold set to "-1".
Read both data/orders and data/products Parquet files.
Revenue calculation (quantity * unit_price).
Join: orders.join(products, "product_id").
Explain: joined.explain(True).
Category aggregation: joined.groupBy("category") with F.sum("revenue").
Show results ordered by descending revenue, then spark.stop().

Your pipeline now joins a million orders with a thousand products and shows revenue by category. But look at that execution plan one more time. Spark is shuffling both tables just to match a 1,000-row lookup table. In the next step, you'll fix this with a broadcast join that eliminates the shuffle entirely.

Fix the Join with Broadcast and AQE

In the last step, you diagnosed the problem. The execution plan revealed a SortMergeJoin with two Exchange nodes. Spark shuffled both the million-row orders table and the tiny products table across partitions just to join them.

The products table is only 1,000 rows. Sending a copy of it to every executor eliminates the shuffle on the orders side entirely. This is called a broadcast join. In this step, you will apply that fix and re-enable Adaptive Query Execution (AQE) so Spark can handle remaining optimizations at runtime.

In this step, get ready to:

Re-enable AQE and skew join handling in the Spark config.
Apply a broadcast join to eliminate the shuffle on the large table.

Re-enable AQE and skew join handling

In Step 3, you deliberately disabled AQE and auto-broadcast to expose the raw SortMergeJoin plan. Now it is time to turn those optimizations back on.

In pipeline.py, find the two .config() lines inside the SparkSession builder.
Find these lines:

    .config("spark.sql.adaptive.enabled", "false") \
    .config("spark.sql.autoBroadcastJoinThreshold", "-1") \

Replace those two lines with the following:

    .config("spark.sql.adaptive.enabled", "true") \
    .config("spark.sql.adaptive.skewJoin.enabled", "true") \

What changed and why?

You made two changes to the Spark config.

spark.sql.adaptive.enabled is now true. AQE re-optimizes the execution plan at runtime using actual data sizes, instead of relying on pre-run estimates that can be wrong.
spark.sql.adaptive.skewJoin.enabled replaces the old autoBroadcastJoinThreshold line. When AQE detects that one partition is much larger than the rest (skew), it automatically splits that partition so no single task becomes a straggler.

Apply the broadcast join and compare plans

Re-enabling AQE helps Spark optimize at runtime, but you can do even better by telling Spark explicitly that the products table is small enough to broadcast.

In pipeline.py, find this section:

# Join orders with products on product_id
joined = orders.join(products, "product_id")

# Inspect the physical plan - look for SortMergeJoin and Exchange nodes

Replace those four lines with the following:

# Join orders with products using broadcast for the small dimension table
joined = orders.join(F.broadcast(products), "product_id")

# Inspect the optimized plan

What does F.broadcast() do?

F.broadcast(products) tells Spark to send a full copy of the products DataFrame to every executor in the cluster. Each executor can then join its local chunk of orders against the local copy of products without any network transfer.

This only works when one side of the join is small enough to fit in memory on each executor. At 1,000 rows, the products table qualifies easily. The million-row orders table stays put.

Before you run the updated pipeline, think about what you expect the execution plan to look like now. Will it still show SortMergeJoin with two Exchange nodes, or something different?

Save pipeline.py.
Run your updated pipeline by running this command:

python pipeline.py

Look at the physical plan output. Instead of SortMergeJoin with two Exchange hashpartitioning nodes, you should now see BroadcastHashJoin with a BroadcastExchange. The Exchange on the orders side is completely gone.

How do I read this new plan?

BroadcastHashJoin means Spark is using a hash-based join where one side was broadcast. No sorting step is needed.
BroadcastExchange appears only on the products side. Spark sends the small table to every executor. That is the only data that moves over the network.
The orders side has no Exchange node at all. The million-row table stays where it is. This is the performance win.

The category revenue results are the same as before. The plan changed, not the data.

Still seeing SortMergeJoin in the plan?

Check that you changed spark.sql.adaptive.enabled from "false" to "true". If it is still "false", AQE cannot help optimize the plan.
Make sure the join line reads F.broadcast(products) with the parentheses wrapping products. A missing closing parenthesis will cause a syntax error instead.
Confirm you removed the autoBroadcastJoinThreshold line entirely. If it is still set to -1, Spark will refuse to auto-broadcast even when AQE is on.
Why is my plan still showing SortMergeJoin after adding F.broadcast()?

Why this matters in interviews

Interviewers testing Spark knowledge almost always ask about join strategies. Being able to say "I switched from SortMergeJoin to BroadcastHashJoin by wrapping the small table in F.broadcast(), which eliminated the shuffle on the fact table" is a concrete, specific answer that demonstrates hands-on experience.

✔️ Awesome, I've got everything!

Great. Double check you have saved pipeline.py.

ⓧ I'd like to double check the full code

Scroll up through your pipeline.py and confirm each section is present in this order:

Imports: SparkSession and functions as F.
SparkSession builder: spark.sql.adaptive.enabled set to "true" and spark.sql.adaptive.skewJoin.enabled set to "true". No autoBroadcastJoinThreshold line.
Read data: orders and products read from data/orders and data/products.
Revenue calculation: quantity * unit_price stored in a revenue column.
Broadcast join: orders.join(F.broadcast(products), "product_id").
Explain + aggregate: joined.explain(True) followed by a groupBy("category") aggregation with total_revenue.
Show + stop: category_revenue.orderBy(...).show() and spark.stop().

Your join is now fast. But 30% of your data is concentrated on just 5 products. Next up, you will handle that skew and write production-ready partitioned output that downstream consumers can query efficiently.

Handle Skew and Write Optimized Output

In the last step, you replaced the slow SortMergeJoin with a BroadcastHashJoin. The execution plan is clean and the shuffle on the orders table is gone.

There is still a hidden problem. 30% of your order data is concentrated on just 5 product IDs. At scale, this skew creates straggler tasks that hold up the entire job. On top of that, writing results as a single blob means downstream consumers have to scan everything. This step enables Spark's runtime skew handling and writes production-ready partitioned output.

In this step, get ready to:

Enable AQE partition coalescing and enrich the aggregation.
Write partitioned Parquet output organized by category.
Verify the partition directory structure.

Enable AQE coalescing and enrich the aggregation

Adaptive Query Execution can automatically right-size shuffle partitions at runtime, merging many small partitions into fewer, larger ones. You need to enable its coalescing feature explicitly.

In pipeline.py, find the SparkSession builder. Locate these two lines:

    .config("spark.sql.adaptive.skewJoin.enabled", "true") \
    .getOrCreate()

Add a new .config() line between them for partition coalescing, so the block now looks like this:

    .config("spark.sql.adaptive.skewJoin.enabled", "true") \
    .config("spark.sql.adaptive.coalescePartitions.enabled", "true") \
    .getOrCreate()

What is partition coalescing?

Spark defaults to 200 shuffle partitions. For a 1-million-row local job, most of those partitions end up nearly empty. Coalescing lets AQE merge them at runtime based on actual data size, so fewer tasks run with more data each.

Now scroll down to the aggregation block. Find this code:

category_revenue = joined.groupBy("category").agg(
    F.sum("revenue").alias("total_revenue")
)

Replace it with an enriched version that also counts orders and calculates average revenue:

category_revenue = joined.groupBy("category").agg(
    F.sum("revenue").alias("total_revenue"),
    F.count("order_id").alias("order_count"),
    F.avg("revenue").alias("avg_order_revenue")
)

Before you run this, think about what the output table will look like now. How many columns do you expect to see?

Save pipeline.py and run the script to see the enriched output:

python pipeline.py

The revenue table now shows four columns: category, total_revenue, order_count, and avg_order_revenue. The physical plan still shows BroadcastHashJoin from the previous step.

Seeing a syntax error?

Check that every .agg() argument except the last one ends with a comma. A missing comma between F.sum() and F.count() is the most common cause.

Also verify the backslash continuation line (\) on the new .config() line you added to the SparkSession builder. It must appear before the line break, with no space after it.

Help me debug a syntax error in my PySpark pipeline.py aggregation

Write optimized output to Parquet

Your pipeline prints results to the terminal, but production pipelines write data to files that downstream systems can query. Parquet files partitioned by a key column let those downstream systems read only the slices they need.

In pipeline.py, find the spark.stop() line at the bottom of the file. Add the following code above it:

# Write optimized output partitioned by category
joined.repartition("category") \
    .write \
    .mode("overwrite") \
    .partitionBy("category") \
    .parquet("output/orders_by_category")

# Write summary
category_revenue.coalesce(1) \
    .write \
    .mode("overwrite") \
    .parquet("output/category_revenue")

print("Output written to output/orders_by_category and output/category_revenue")

What does this code do?

.repartition("category") reshuffles the joined data so all rows for the same category land in the same partition. .partitionBy("category") then writes each partition as its own subdirectory on disk (e.g. category=Electronics/).

.coalesce(1) merges all summary partitions into a single output file. This is appropriate here because the summary is only 5 rows (one per category). You would never use coalesce(1) on a large dataset.

Save pipeline.py and run the script again:

python pipeline.py

After the revenue table prints, you should see the line Output written to output/orders_by_category and output/category_revenue. The script now writes files to disk before stopping.

Seeing a permission or path error?

Make sure you are running the script from the same directory that contains your data/ folder. The output/ directory is created automatically. If you see mode("overwrite") errors, check that no other process has a lock on the output directory.

Help me debug a file write error in my PySpark pipeline

Inspect the partition structure

Before you check the output, predict what subdirectories you expect to see inside output/orders_by_category/. How many categories did your data have?

List the output directory to confirm the partition structure by running:

ls output/orders_by_category/

You should see five subdirectories, one for each category: category=Books, category=Clothing, category=Electronics, category=Home, and category=Sports.

Why partition by category?

A downstream query that only needs Electronics data can now read just the category=Electronics/ subdirectory. Without partitioning, it would have to scan all 1 million rows and filter in memory. This is how production data lakes are organized for efficient queries.

Confirm the summary output also exists by running:

ls output/category_revenue/

You should see a single Parquet file (plus a _SUCCESS marker). That single file holds your 5-row category summary.

✔️ Awesome, I've got everything!

Great. Double check you have saved pipeline.py and that the output/ directory contains both orders_by_category/ and category_revenue/.

ⓧ I'd like to double check the full code

The full pipeline.py is longer than 30 lines, so here is a checklist instead. Scroll through your file and confirm each section is present, in this order:

Imports: SparkSession and functions as F.
SparkSession builder: Three .config() lines for adaptive.enabled, skewJoin.enabled, and coalescePartitions.enabled.
Read data: orders and products from Parquet.
Revenue calculation: quantity * unit_price.
Broadcast join: F.broadcast(products) in the join call.
Explain: joined.explain(True).
Aggregation: F.sum("revenue"), F.count("order_id"), and F.avg("revenue").
Show results: orderBy(F.col("total_revenue").desc()).show().
Partitioned write: .repartition("category") followed by .partitionBy("category") to output/orders_by_category.
Summary write: .coalesce(1) to output/category_revenue.
Print confirmation and spark.stop() at the end.

Your pipeline is complete. It reads raw data, joins it with a broadcast, aggregates revenue by category, and writes production-ready partitioned output. Next up, tackle the Secret Mission and deploy this pipeline to the cloud.

Secret mission

Deploy to AWS EMR Serverless

You have a working local pipeline. Now take it to the cloud. Create an S3 bucket, configure an IAM role, and submit your optimized PySpark job to AWS EMR Serverless. This is the difference between understanding Spark and operating it in production.

Clean Up Your Resources

This project runs entirely on your local machine, so there are no ongoing costs. Decide whether to keep your files to experiment further, or clean them up to free disk space.

Resources you used:

Project scripts: generate_data.py, pipeline.py, and pipeline_emr.py (Secret Mission only).
data/ directory containing generated Parquet files for orders and products.
output/ directory containing pipeline results (partitioned Parquet and summary).
PySpark 4.2.0 Python package installed via pip.

What about AWS resources?

If you completed the Secret Mission, your AWS resources (EMR Serverless application, S3 bucket, and IAM role) were already cleaned up in that section. If you skipped the Secret Mission, no cloud resources were ever created.

Keep everything running

No action needed. Choose this if you want to keep experimenting with Spark configurations, data sizes, or aggregation patterns.

Your local files take minimal disk space. The data/ directory with 1 million rows of Parquet data is typically under 50 MB.
Java 17 and PySpark remain installed on your machine. Both are useful for future Spark or JVM-based projects.
Try modifying generate_data.py to change the skew ratio or increase the number of orders to see how execution plans change at different scales.

Pause - I'll come back to this later

Free up disk space by removing generated data while keeping your scripts for later. You can regenerate the data at any time.

Delete the data/ and output/ directories to reclaim disk space. Your scripts can regenerate them.
Remove the data and output directories by running:

rm -rf data/ output/

Keep generate_data.py and pipeline.py so you can pick up where you left off. Run python generate_data.py to recreate the data when you return.
Java 17 and PySpark stay installed. No action needed for those.

Delete - I don't want to use this again

Remove all project files and uninstall PySpark to start completely fresh.

Remove all project files and data by running the commands for your operating system.

macOS

Remove the data directories and project scripts by running:

rm -rf data/ output/
rm generate_data.py pipeline.py

If you completed the Secret Mission, also remove the EMR script by running:

rm pipeline_emr.py

Windows

Remove the data directories and project scripts by running:

rmdir /s /q data output
del generate_data.py pipeline.py

If you completed the Secret Mission, also remove the EMR script by running:

del pipeline_emr.py

Uninstall PySpark by running:

pip uninstall pyspark -y

Java 17 remains installed on your machine. Other tools and projects may depend on it, so only remove it if you are sure nothing else needs it.

Nice Work!

Nice work! You just built, broke, diagnosed, and fixed a production-grade PySpark batch pipeline from scratch.

You've learned how to:

Built a PySpark batch pipeline that generates, joins, and aggregates 1 million e-commerce records, then writes the results as production-ready partitioned Parquet output organized by category.
Diagnosed a slow SortMergeJoin by reading Spark execution plans and identifying Exchange (shuffle) nodes that would kill performance at scale.
Fixed the pipeline using F.broadcast() to eliminate the shuffle, re-enabled Adaptive Query Execution with skew handling, and confirmed the plan changed from SortMergeJoin to BroadcastHashJoin.
Completed an optional Secret Mission to push your skills further and prove you can operate Spark in a cloud environment.

Ready to quiz yourself?