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

Port a pandas aggregation to the Pandas API on Spark

You have an old pandas snippet that summarizes order lines per company. It works in a notebook, but only on data that fits in memory. Port it to the Pandas API on Spark (pyspark.pandas), so that it runs distributed and lazily. Keep the pandas ergonomics.

Implement line_summary(order_items) in solution.py. order_items is a Spark DataFrame with company_id, order_item_id, order_id, product_code, qty, line_amount. Return a Spark DataFrame grouped by company_id with these columns:

column	meaning
`company_id`	the group key
`n_lines`	number of order-line rows in the company
`total_line_amount`	sum of `line_amount` in the company

Do the aggregation through the Pandas API on Spark (df.pandas_api() → pandas-style groupby → .to_spark()). Do not use collect() to move data to the driver. The test compares against a DuckDB oracle. It ignores the order of rows and columns, and floats use tolerance.

Note: the Pandas API on Spark is distributed pandas. It is the right tool when the data is too big for one machine. For small data, real pandas or DuckDB is simpler. This is a porting exercise, so the dataset is small on purpose.

Starter code

solution.py

"""SPK-09 / 03-pandas-api — starter. See README.md.

Port the pandas-style aggregation to the Pandas API on Spark and return a Spark
DataFrame with company_id, n_lines, total_line_amount.
"""


def line_summary(order_items):
    # TODO: use order_items.pandas_api(), a pandas-style groupby, then .to_spark().
    raise NotImplementedError

Hints

Hint 1

Convert the Spark DataFrame to a Pandas-API-on-Spark frame with df.pandas_api() (or df.to_pandas_on_spark()). Then write the familiar pandas groupby, and convert the result back with .to_spark().

Hint 2

Group by company_id; n_lines = size of each group; total_line_amount = sum of line_amount. Reset the index so company_id is a column, then rename to the required names.

Hint 3

Return a Spark DataFrame (call .to_spark()) with columns company_id, n_lines, total_line_amount. Row/column order are not checked; floats use tolerance.

Run it locally

uv run de start SPK-09/03-pandas-api
# edit workspace/SPK-09/03-pandas-api/
uv run de test