DE Mastery GitHub

Apache Spark · SPK-04

Spark SQL and data sources

From your world. dbt {{ ref() }} and marts vs. temp views and SQL on files

← All exercises

Write partitioned output with the right save mode

A nightly job writes the orders table to a lake path. The data is partitioned by company, so downstream reads can prune. The job must be idempotent. If you run it twice, exactly one copy of the data must remain, not two.

Implement write_partitioned(df, out_path) in solution.py. It must write df to out_path as Parquet. The output must meet these conditions:

  1. The output is partitioned by company_id (one company_id=.../ directory for each company).
  2. The save mode is overwrite. Then, when you run the job again, it replaces the previous output. It does not append a second copy.

The function writes data. It does not need to return anything.

The test does these steps:

  • It calls write_partitioned once, reads the path back, and compares the result to the input. This is an independent check that no rows are lost or duplicated.
  • It confirms that the company_id=... partition directories exist.
  • It calls write_partitioned again and confirms that the row count is the same. This is where append would fail, because it would double the data.

The partition column matters. company_id has low cardinality and queries filter on it, so it prunes well. Do not partition by a high-cardinality key like order_id. It would create a flood of tiny files.

Starter code

solution.py
"""SPK-04 / 02-write-partitioned — starter. See README.md."""


def write_partitioned(df, out_path):
    # TODO: write df to out_path as Parquet, partitioned by company_id, mode overwrite.
    raise NotImplementedError

Hints

Hint 1

df.write.mode('overwrite').partitionBy('company_id').parquet(out_path) writes one directory per company and replaces any previous output.

Hint 2

Use overwrite, NOT append: a refresh job that appends stacks a second copy and doubles every metric. The test runs your writer twice and checks the row count is stable.

Hint 3

partitionBy writes the partition column as a directory name (company_id=...), so the written files no longer contain that column physically — but reading the path back restores it as a normal column.

Run it locally
uv run de start SPK-04/02-write-partitioned
# edit workspace/SPK-04/02-write-partitioned/
uv run de test