Handle Data Skew from High-Frequency Vehicles (PySpark)

PYSPARK coding challenge · Difficulty: hard · Topic: Performance Optimization · +150 XP

Scenario

You are a Data Engineer at Uber. Every vehicle on the platform streams GPS and

sensor telemetry into a single table. Most vehicles emit a comparable number of

events, but a small set of high-frequency vehicles — airport shuttles, long-haul

fleet cars, always-on test rigs — emit orders of magnitude more.

When you aggregate by vehicle_id, every row for one vehicle hashes to the **same

partition**. A handful of reducers receive tens of millions of rows while the rest

receive a few thousand. The stage finishes only when its slowest task finishes, so

the whole job waits on those few partitions. That is data skew.

Your task is to compute daily per-vehicle statistics in a way that spreads the

skewed keys across partitions, while returning a result identical to a plain

aggregation.

Table: telemetry

| Column | Type | Description |
| --- | --- | --- |
| `vehicle_id` | string | Unique vehicle identifier |
| `event_time` | timestamp | When the event occurred |
| `speed` | double | Speed in km/h |
| `distance_km` | double | Distance covered during the event |

In the sample below V001 is the high-frequency vehicle: it accounts for **18 of

25 rows (72%)**, which is the skew you must handle. In production the same shape

holds at a scale of billions of rows.

Task

Using the PySpark DataFrame API, return one row per vehicle per day with:

| Column | Meaning |
| --- | --- |
| `vehicle_id` | the vehicle |
| `event_date` | date part of `event_time` |
| `total_events` | number of telemetry events |
| `total_distance_km` | sum of `distance_km`, rounded to 2 decimals |
| `avg_speed` | mean of `speed`, rounded to 2 decimals |
| `max_speed` | maximum `speed` |

Order by vehicle_id, then event_date. Assign the result to df_result and call

df_result.show().

Required Approach

A plain groupBy("vehicle_id", "event_date") is not an acceptable answer here.

Use key salting:

1. Add a salt column — a bounded random or hashed bucket, e.g. 0..15.

2. Stage 1: aggregate by vehicle_id, event_date, salt. The skewed vehicle's

rows now spread across 16 partial groups instead of one.

3. Stage 2: aggregate the partial results by vehicle_id, event_date to collapse

the salt away.

The trap: you cannot average an average

avg(avg(speed)) is wrong whenever the salt buckets hold different numbers of

rows — which is exactly the case with random salting.

Carry sum(speed) and count(*) through stage 1, then divide once in stage 2:

`

avg_speed = sum(partial_speed_sum) / sum(partial_count)

`

sum, count and max are safe to combine this way because they are associative.

avg is not. This is the single most common way a salted aggregation silently

returns wrong numbers.

Example Input

| vehicle_id | event_time | speed | distance_km |
| --- | --- | --- | --- |
| V001 | 2026-08-01 08:01:00 | 60.0 | 1.25 |
| V001 | 2026-08-01 08:02:00 | 64.0 | 1.50 |
| ... | (18 rows for V001 in total) | ... | ... |
| V002 | 2026-08-01 08:01:00 | 45.0 | 0.50 |
| V003 | 2026-08-01 08:10:00 | 52.0 | 0.75 |

Expected Output

| vehicle_id | event_date | total_events | total_distance_km | avg_speed | max_speed |
| --- | --- | --- | --- | --- | --- |
| V001 | 2026-08-01 | 12 | 13.5 | 62.0 | 66.0 |
| V001 | 2026-08-02 | 6 | 9.0 | 71.33 | 74.0 |
| V002 | 2026-08-01 | 3 | 3.0 | 50.0 | 55.0 |
| V002 | 2026-08-02 | 2 | 1.0 | 43.0 | 46.0 |
| V003 | 2026-08-01 | 2 | 2.0 | 55.0 | 58.0 |

Constraints

• telemetry may contain billions of rows; a small percentage of vehicles may

account for more than 50% of them.

• Do not call collect() or toPandas() on the full dataset.
• Do not use a Python UDF — it forces serialization to the Python worker and

defeats the optimization you are demonstrating.

• The result must be exactly equal to the plain aggregation.

Follow-up (discuss in your head, not in code)

• Why does groupBy("vehicle_id") skew, and where exactly does the job stall?
• How does salting change the partition distribution for V001?
• Why is the two-stage result provably equal to the one-stage result?
• When would you reach for AQE skew-join optimization (`spark.sql.adaptive.

skewJoin.enabled`) or a broadcast join instead of salting?

> F (pyspark.sql.functions) and the telemetry DataFrame are already available.

What this PYSPARK challenge teaches you

“Handle Data Skew from High-Frequency Vehicles (PySpark)” is a hard-level PYSPARK challenge focused on Performance Optimization. Working through it gives you hands-on practice with data-skew, salting, aggregation, partitioning, shuffle, performance — 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

• data-skew
• salting
• aggregation
• partitioning
• shuffle
• performance
• aqe
• uber

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. Plain groupBy('vehicle_id') sends every row for V001 to one reducer. Add a bounded salt column so those rows spread across many partitions.
2. Stage 1: groupBy('vehicle_id','event_date','salt'). Stage 2: groupBy('vehicle_id','event_date') to collapse the salt away.
3. Do not carry avg through stage 1. Carry sum(speed) and count(*), then divide once in stage 2 -- avg of avgs is wrong when buckets hold different row counts.
4. sum, count and max are associative so the two-stage result equals the one-stage result. Round total_distance_km and avg_speed to 2 decimals.

Where this comes up

Variations of this problem have been reported in interviews at Uber. Interviewers use it to check whether you can express the logic cleanly and reason about correctness on edge cases such as ties, nulls and empty groups.

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

• Fix the Broken Pipeline
• Optimize the 100M-Row Join
• Parse Apache Logs with Regex
• Optimize Small DataFrame Join with Broadcast
• Optimize Average Rating Calculation for Products
• Deduplicate and Aggregate User Actions with Latest Session

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 Performance Optimization.

Solve this challenge free on PySpark.in