DE Mastery GitHub

Apache Spark · SPK-09

UDFs and the Pandas API

From your world. dbt Python models, and the pandas you reach for out of habit

← All exercises

Vectorize a currency-code cleaner

The orders feed has dirty currency codes. It has lowercase usd, blank strings, and the bogus XXX. These are mixed with the real codes USD/EUR/GBP/JPY. A teammate cleans them with a per-row Python UDF (BatchEvalPython). Produce the same cleaned output with no such fence.

Implement clean_currency(orders) in solution.py. You get the raw orders DataFrame. Return a DataFrame with the original order_id plus these columns:

column meaning
currency_clean canonical form: trimmed and upper-cased; a blank code becomes NULL
is_valid boolean: currency_clean is one of USD, EUR, GBP, JPY (NULL is not valid)

Rules:

  • Canonicalize before you judge. "usd" → "USD" is valid. " eur " → "EUR" is valid.
  • A blank/empty currency becomes currency_clean = NULL and is_valid = false.
  • "XXX" canonicalizes to "XXX" and is is_valid = false (it is not an ISO code that we accept).
  • The physical plan must contain no BatchEvalPython. Native functions or a pandas/Arrow UDF are fine. A plain Python udf(...) is not.

The test checks the values against a DuckDB oracle. It also asserts that the plan has no BatchEvalPython.

Starter code

solution.py
"""SPK-09 / 02-normalize-codes — starter.

Works, but cleans the currency with a per-row Python UDF -> BatchEvalPython. Rewrite
so the plan has no BatchEvalPython (see README.md).
"""

from pyspark.sql import functions as F
from pyspark.sql.functions import udf

_VALID = {"USD", "EUR", "GBP", "JPY"}


def clean_currency(orders):
    def canon(code):
        if code is None:
            return None
        c = code.strip().upper()
        return None if c == "" else c

    def valid(code):
        c = canon(code)
        return c in _VALID

    canon_udf = udf(canon, "string")
    valid_udf = udf(valid, "boolean")
    return orders.select(
        "order_id",
        canon_udf("currency").alias("currency_clean"),
        valid_udf("currency").alias("is_valid"),
    )

Hints

Hint 1

The starter normalizes currency with a per-row Python UDF (BatchEvalPython). Everything it does — upper-case, trim, map empty to NULL, flag non-ISO codes — has a native equivalent in pyspark.sql.functions.

Hint 2

Use F.upper(F.trim(col)) to canonicalize; F.when(code == '' , None) (or nullif) to turn blanks into NULL; and F.col('code').isin(valid_set) to decide is_valid.

Hint 3

Return currency_clean (canonical upper/trim, NULL for blank) and is_valid (boolean: canonical code is one of USD/EUR/GBP/JPY). A pandas_udf over a Series is an acceptable alternative; a plain Python udf is not.

Run it locally
uv run de start SPK-09/02-normalize-codes
# edit workspace/SPK-09/02-normalize-codes/
uv run de test