Top N Per Group

PySpark data engineering interview problem. Difficulty: intermediate. Pattern: Window Functions. About 16 minutes. Part of the Pro drill bank.

Top 3 amounts per user_id via row_number. Treat this as a production helper: match the contracted return shape, including empty and duplicate inputs.

A very common question is: for each group, give me the top few rows. For example, the 3 biggest amounts for every user. This is called top N per group. For each user_id, sort the rows by amount from high to low and number them with row_number(). Keep only the first 3 rows of each user. Return user_id, amount and the row number as rn, sorted by user_id and rn. The editor creates the table df and ends with result = df. Change the last line to build your answer.

Requirements

• Use row_number() over a window partitioned by user_id and ordered by amount descending.
• Keep only the rows where rn is 3 or less.
• Return user_id, amount, rn, sorted by user_id, rn.
• Store the DataFrame in result.

Constraints

• Keep the rn column in the output. Do not drop it.
• A user with fewer than 3 rows keeps all their rows.
• Sort the amounts from high to low, so row number 1 is the biggest amount.

Examples

Input: amounts per user Output: user_id | amount | rn 1 | 40.0 | 1 1 | 30.0 | 2 1 | 20.0 | 3 2 | 50.0 | 1 2 | 5.0 | 2 User 1 drops the 10.0 fourth place; user 2 has only 2 rows.

Topics: lakebench, pyspark, row_number, top-n.

More PySpark interview questions · All interview problems · Learn data engineering