Optimize Join Operation Using Broadcast in AQE

PYSPARK coding challenge · Difficulty: hard · Topic: Adaptive Query Execution · +200 XP

Problem

You are given two DataFrames: customers and orders. The orders DataFrame is significantly larger than the customers DataFrame. Your task is to join these two DataFrames on the customer_id column using Spark's Adaptive Query Execution (AQE) to efficiently handle the skewed join operation. Use broadcasting for the smaller DataFrame to optimize the query.

Example Input - `customers`

| customer_id | name |
| --- | --- |
| 1 | Alice |
| 2 | Bob |
| 3 | Cathy |

Example Input - `orders`

| order_id | customer_id | amount |
| --- | --- | --- |
| 101 | 1 | 250 |
| 102 | 1 | 300 |
| 103 | 2 | 150 |
| 104 | 3 | 400 |
| 105 | 3 | 500 |
| 106 | 3 | 450 |

Expected Output

| customer_id | order_id | amount | name |
| --- | --- | --- | --- |
| 1 | 101 | 250 | Alice |
| 1 | 102 | 300 | Alice |
| 2 | 103 | 150 | Bob |
| 3 | 104 | 400 | Cathy |
| 3 | 105 | 500 | Cathy |
| 3 | 106 | 450 | Cathy |

Notes

What this PYSPARK challenge teaches you

“Optimize Join Operation Using Broadcast in AQE” is a hard-level PYSPARK challenge focused on Adaptive Query Execution. Working through it gives you hands-on practice with broadcast, join, AQE — the kind of transformation you are asked to write in real data engineering work and in technical interviews. You can solve it directly in the browser: the dataset is pre-loaded, so you write the query or DataFrame code, run it, and compare your output against the expected result immediately.

Concepts covered

How to approach it

If you get stuck, work through these steps in order before looking at a full solution — each one narrows the problem down:

  1. 1 ≤ number of rows in customers ≤ 10,000
  2. 1 ≤ number of rows in orders ≤ 1,000,000
  3. Broadcast only if the DataFrame has less than 10,000 rows
  4. Enable AQE for better join strategy optimization

How to practise it on PySpark.in

Open the challenge, write your PySpark code in the editor and press Run to execute it against the sample dataset. Submitting checks your output against every test case, including hidden ones, so you find out straight away whether your logic holds up. You can retry as often as you like, and each solved challenge adds to your XP.

Related PYSPARK challenges

Frequently asked questions

Do I need to install Spark or a database to solve this?

No. The PYSPARK environment runs in your browser with the sample data already loaded, so there is nothing to install or configure.

Is this challenge free?

Yes - the problem, the sample dataset, the hints and unlimited test runs are free.

What level is it?

It is rated hard and covers Adaptive Query Execution.

Solve this challenge free on PySpark.in