UDF Performance Trap

PySpark data engineering interview problem. Difficulty: intermediate. Pattern: Explain Plans. About 16 minutes. Free to practice.

A colleague banded the order amounts with a Python UDF. On a small test it worked, but on the full data the job crawled. A UDF makes Spark send every value to a Python process and back, and Spark cannot optimize what happens inside it. Your task is to write the same logic with built-in functions.

Put each amount into one of three bands: Low for amounts below 50, Medium for 50 up to and including 200, and High above 200. Use F.when(...).otherwise(...) and not a UDF, so Spark can run the logic inside its own engine. 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

• Create band with built-in when/otherwise.
• Use the same rules as before: below 50 is Low, up to 200 is Medium, above is High.
• Return user_id, amount, band, sorted by user_id.
• Store the DataFrame in result.

Constraints

• Do not use F.udf or pandas_udf.
• All 6 rows stay in the output.
• The band names are exactly Low, Medium and High.

Examples

Input: amounts Output: user_id | amount | band 1 | 20.0 | Low 4 | 250.0 | High 6 | 200.0 | Medium Same labels without a UDF.

Topics: lakebench, pyspark, when, performance.

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