Clean, split, merge, and aggregate with pandas
Company: Uber
Role: Data Scientist
Category: Data Manipulation (SQL/Python)
Difficulty: medium
Interview Round: Technical Screen
Given two CSVs, use pandas to clean, split strings, merge, and aggregate.
drivers.csv
driver_id,name,signup_city
D1,Jane Doe,SF
D2,Mark S,NYC
D3,Adam L,LA
trips.csv
trip_id,driver_id,ts_utc,route,fare_usd
T1,D1,2019-01-01T08:00:00Z,SF-CA|SFO,17.52
T2,D1,2019-01-02T09:00:00Z,SF-CA|DAL,4.40
T3,D2,2019-01-02T10:00:00Z,,
T4,D3,2019-01-03T11:00:00Z,LA-CA|LAX,8.00
Tasks:
1) Load both files into DataFrames; show head(2) and tail(1) of trips to verify ingest.
2) Drop rows in trips with missing fare_usd or missing/empty route (after stripping whitespace). Ensure fare_usd is numeric.
3) Split the route column on '|' into origin and destination columns; trim whitespace. If the split yields fewer than 2 tokens, drop those rows.
4) Merge the cleaned trips with drivers on driver_id (left join from trips) and keep only rows with a matching driver.
5) Produce a per-driver summary with columns: driver_id, name, trips_count, avg_fare_usd (rounded to 2 decimals), last_3_trips_avg (average of the last 3 trips per driver ordered by ts_utc; if <3 trips, average over available).
6) Return the top 2 drivers by avg_fare_usd DESC (break ties by trips_count DESC, then name ASC) and print the final DataFrame schema (dtypes) to confirm transformations.
Explicitly use: head, tail, dropna, str.split, merge, groupby, sort_values, and rolling/agg as appropriate.
Overview: This question evaluates proficiency with pandas data manipulation and cleaning, including handling missing values, trimming and splitting strings, numeric type conversion, merges, time-aware ordering, and aggregation for per-entity summaries.
You are given two tables, `drivers` and `trips`, that store information about ride-hailing drivers and the trips they completed.
Write a **single PostgreSQL `SELECT`** query (CTEs and window functions are allowed) that cleans, splits, joins, and aggregates the data as follows, and returns the final result.
**1) Clean the trips.** Keep only rows from `trips` where:
- `fare_usd` is not NULL, **and**
- `route` is not NULL and is non-empty after trimming surrounding whitespace.
**2) Split the route.** The `route` column holds strings like `'SF-CA|SFO'`. Split each route on the `'|'` character into `origin` (the part before `'|'`) and `destination` (the part after `'|'`). If a `route` value does **not** contain a `'|'`, drop that trip.
**3) Join to drivers.** Inner-join the cleaned trips to `drivers` on `driver_id`, keeping only trips whose `driver_id` matches an existing driver.
**4) Aggregate per driver.** For each driver that still has at least one trip, compute:
- `trips_count` — the total number of remaining trips for that driver.
- `avg_fare_usd` — the average of `fare_usd` over that driver's remaining trips, rounded to 2 decimal places.
- `last_3_trips_avg` — the average `fare_usd` over that driver's **3 most recent trips** by `ts_utc` (newest 3). If the driver has fewer than 3 trips, average over all of their trips. Compute this with a window function (a rolling average over the current row and the 2 preceding rows when ordered by `ts_utc` ascending), then collapse to the single value belonging to the driver's most recent trip. Round to 2 decimal places.
**5) Rank and limit.** Return only the **top 2** drivers, ordered by:
- `avg_fare_usd` descending, then
- `trips_count` descending (tie-break), then
- `name` ascending (final tie-break).
The result must have **one row per driver** and exactly these columns, in order: `driver_id`, `name`, `trips_count`, `avg_fare_usd`, `last_3_trips_avg`.
Tables
drivers(driver_id VARCHAR(10), name VARCHAR(100), signup_city VARCHAR(50))
trips(trip_id VARCHAR(10), driver_id VARCHAR(10), ts_utc TIMESTAMP, route VARCHAR(100), fare_usd DECIMAL(10,2))
Hints
- In PostgreSQL, split a delimited string with `SPLIT_PART(route, '|', 1)` and `SPLIT_PART(route, '|', 2)` — there is no `SUBSTRING_INDEX`.
- Filter out unusable routes by requiring both non-NULL/non-blank and the presence of the delimiter, e.g. `route LIKE '%|%'`.