Simple UDF

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

Band amounts Low/Medium/High with when (no Python UDF). Treat this as a production helper: match the contracted return shape, including empty and duplicate inputs.

In interviews, people are often asked to write a UDF (user defined function), a Python function that Spark calls for each row. UDFs are slow, because Spark must move each value to Python and back. When the logic is simple, built-in functions are almost always better. Here, put each amount into a band: Low for amounts below 50, Medium for amounts from 50 up to and including 200, and High for amounts above 200. Use F.when(...).otherwise(...) and not a UDF. Return user_id, amount and band, sorted by user_id. The editor creates the table df and ends with result = df. Change the last line to build your answer.

Requirements

• Add a band column with the values Low, Medium and High.
• Low is below 50, Medium is 50 to 200 including both ends, and High is above 200.
• Use when/otherwise, not a UDF.
• Return user_id, amount, band, sorted by user_id, in result.

Constraints

• The band names are Low, Medium and High, with a capital first letter.
• Chain the conditions from the lowest to the highest. when checks them in order.
• 50 and 200 both belong to Medium.

Examples

Input: amounts Output: user_id | amount | band 1 | 20.0 | Low 2 | 80.0 | Medium 3 | 150.0 | Medium 4 | 250.0 | High 5 | 50.0 | Medium 6 | 200.0 | Medium Inclusive Medium on both endpoints.

Topics: lakebench, pyspark, when, udf-reframe.

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