Pandas UDF (Vectorized)

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

Same banding as problem 33; pandas_udf is unsupported in the shim. Treat this as a production helper: match the contracted return shape, including empty and duplicate inputs.

A pandas UDF (or vectorized UDF) is faster than a normal UDF, because it processes many rows at once. But the simplest and fastest solution is often to avoid UDFs completely and use Spark's built-in functions. In this editor, pandas_udf is not available. Solve the banding problem from the previous exercise without it, using when and otherwise. Put each amount into a band: Low below 50, Medium from 50 to 200 including both ends, and High above 200. 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

• Do not use pandas_udf or any UDF.
• Create a band column with Low, Medium and High, using when/otherwise.
• Return user_id, amount, band, sorted by user_id.
• Store the DataFrame in result.

Constraints

• The rules are the same as the previous exercise: < 50 is Low, up to 200 is Medium, above is High.
• The band names are written with a capital first letter.
• Keep all rows. Banding adds a column and removes nothing.

Examples

Input: amounts Output: user_id | amount | band 1 | 20.0 | Low 4 | 250.0 | High Native expressions avoid Python row workers.

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

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