Partitioning, bucketing, and the small-file problem

6 h
requires SPK-06

From your world. Redshift DISTKEY / SORTKEY; Athena/Iceberg partitioning and file sizing

Hit a target output file count

The number of files that you write is the main lever for the small-file problem. If you write too many files, listing takes most of the time. If you write too few, you lose parallelism. In production, you calculate the count from total_bytes / target_size. At lab scale, you target the count directly. This uses the same mechanism: the number of in-memory partitions at write time controls the number of part files.

Task

Implement write_balanced(df, out_dir, num_files). It writes df as Parquet into out_dir so that it produces exactly num_files part files, with the data unchanged. Return the output path as a string.

Constraints

Honor num_files. Do not hard-code coalesce(1).
Produce balanced partitions (a full shuffle to num_files partitions is expected).
Overwrite any existing output so the write is deterministic.

How it is graded

File count: the number of part-*.parquet files in out_dir equals num_files (tested for num_files = 2 and 4).
Data (DuckDB oracle): when you read the written files back, you get exactly the input rows (order-independent).

Starter code

solution.py

"""SPK-07 / 01-target-file-count — starter. See README.md."""


def write_balanced(df, out_dir, num_files):
    # TODO: write df as Parquet to out_dir producing exactly num_files part files,
    # data unchanged. Return the output path (str). Honor num_files; overwrite.
    raise NotImplementedError

Hints

Hint 1

To write exactly N part files, make the DataFrame have exactly N in-memory partitions before writing: df.repartition(num_files).write.parquet(out_dir).

Hint 2

repartition(N) gives N balanced partitions (a full shuffle); each writes one part-*.parquet file. Do NOT hard-code coalesce(1) — honor the num_files argument.

Hint 3

Return the output path (str). The test counts the part-*.parquet files written and reads the data back to compare with the input via the DuckDB oracle.

Run it locally

uv run de start SPK-07/01-target-file-count
# edit workspace/SPK-07/01-target-file-count/
uv run de test