UDFs and the Pandas API

4 h
requires SPK-05

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

Replace a slow per-row FX Python UDF

Finance wants every payment converted to USD. The conversion must use the daily FX rate on the payment date. A teammate shipped a version that works but is slow. It is a Python UDF that looks up the rate for each row. Its physical plan contains a BatchEvalPython node, so every row makes a round trip to a Python worker. Your job is to produce the same numbers without that fence.

Implement to_usd(payments, fx_rates) in solution.py. These are the inputs:

payments(company_id, payment_id, order_id, amount, currency, paid_ts_utc) — amount may be NULL. paid_ts_utc is an ISO UTC timestamp string.
fx_rates(rate_date, currency, usd_per_unit) — one rate per (rate_date, currency) day.

Return a DataFrame with one row per input payment and these columns:

column	meaning
`payment_id`	the payment key
`currency`	the original currency
`amount`	the original amount (may be NULL)
`usd_amount`	`amount * usd_per_unit` for the payment's date; NULL if `amount` or the rate is missing

Rules:

Convert at the rate for the own date of the payment (the calendar date of paid_ts_utc in UTC). Match it against fx_rates.rate_date.
Keep every payment row. A missing rate or a NULL amount must not drop the row (use a left join). It only makes usd_amount NULL.
The physical plan of the result must contain no BatchEvalPython node. You can use native pyspark.sql.functions, a pandas UDF, or an Arrow UDF. You cannot use a plain Python udf(...).

The test checks these three things:

The values match a DuckDB oracle (order ignored, float tolerance).
Every input payment survives.
The plan has no BatchEvalPython.

Starter code

solution.py

"""SPK-09 / 01-kill-the-udf — starter.

This version WORKS numerically but is slow: it looks up the FX rate with a plain
Python UDF, which the planner runs as a BatchEvalPython node (a per-row round trip
to a Python worker). See README.md. Rewrite it so the plan has no BatchEvalPython.
"""

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


def to_usd(payments, fx_rates):
    # Build a {(rate_date, currency): usd_per_unit} lookup on the driver...
    rate_rows = fx_rates.collect()
    rate_map = {(r["rate_date"], r["currency"]): r["usd_per_unit"] for r in rate_rows}

    # ...and apply it per row with a Python UDF -> BatchEvalPython (slow, an optimizer fence).
    def convert(amount, currency, date_str):
        if amount is None:
            return None
        rate = rate_map.get((date_str, currency))
        if rate is None:
            return None
        return amount * rate

    convert_udf = udf(convert, "double")

    return (
        payments
        .withColumn("rate_date", F.to_date(F.to_timestamp("paid_ts_utc")).cast("string"))
        .withColumn("usd_amount", convert_udf("amount", "currency", "rate_date"))
        .select("payment_id", "currency", "amount", "usd_amount")
    )

Hints

Hint 1

The starter multiplies amount by a rate looked up in a Python UDF — that becomes a BatchEvalPython node. The fx_rates table already has one usd_per_unit per (rate_date, currency); join to it instead of looking it up in Python.

Hint 2

Convert paid_ts_utc to a date to match fx_rates.rate_date, join on (currency, rate_date), then usd_amount = amount * usd_per_unit using plain Column arithmetic. No pyspark UDF needed at all.

Hint 3

Keep every payment row even when a rate or amount is NULL: use a LEFT join (broadcast is fine) so you don't silently drop rows. A pandas_udf or Arrow udf is an acceptable alternative, but a plain Python udf is not — the test rejects BatchEvalPython.

Run it locally

uv run de start SPK-09/01-kill-the-udf
# edit workspace/SPK-09/01-kill-the-udf/
uv run de test