Avoid UDF When Possible

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

Prefer when/otherwise over UDFs for simple banding. Treat this as a production helper: match the contracted return shape, including empty and duplicate inputs.

Using a Python UDF to put values in bands works, but it is slow on a cluster, because every value has to be sent from Spark to Python and back. Built-in functions like when are faster, because Spark can run and optimize them inside its own engine. Put each amount into a band using F.when(...).otherwise(...): Low for amounts below 50, Medium for 50 up to and including 200, and High above 200. Return user_id, amount and band, sorted by user_id. Only the DataFrame is graded, but think about why this is better than a UDF. The editor creates the table df and ends with result = df. Change the last line to build your answer.

Requirements

• Create band with F.when and otherwise. Do not use a UDF.
• Low is below 50, Medium is 50 to 200 including both ends, and High is above 200.
• Return user_id, amount, band, sorted by user_id.
• Store the DataFrame in result.

Constraints

• The band names are written exactly Low, Medium and High.
• Keep all 6 rows. Adding a band must not remove any row.
• Check the conditions from the lowest band to the highest.

Examples

Input: amounts Output: user_id | amount | band 1 | 20.0 | Low 2 | 80.0 | Medium 4 | 250.0 | High Built-ins avoid per-row Python.

Topics: lakebench, pyspark, when.

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