partitionBy on Write

Concepts covered: Reading & Writing

partitionBy on the writer turns a column's values into directories. Each distinct value gets a folder named column=value, the rows carrying that value go into files inside it, and the column is dropped from the files themselves because the path already records it.

clicks.write.partitionBy("click_date").parquet("/lake/clicks")

Reading /lake/clicks back, Spark walks the directories, parses each name, and adds click_date as a column again. It infers the type from the names: dates come back as dates and digit strings as integers, so a store code like 00123 written as a partition value reads back as 123. Passing a schema on read, or setting spark.sql.sources.partitionColumnTypeInference.enabled to false, keeps it a string.

How many files a write emits

Each write task opens 1 file for every distinct partition value among the rows it holds. Upstream partitioning decided which rows each task holds without looking at click_date, so a task holding rows from 30 dates writes 30 files. The ceiling is the task count times the value count.

Take a backfill of 30 days of clicks, 30GB in all, written straight after a groupBy that left 200 partitions:

The first row averages 5MB per file, the small-files problem produced by your own write. Hashing on the partition column sends all of a date's rows to 1 task, so every directory gets exactly 1 file:

(clicks
    .repartition("click_date")
    .write
    .partitionBy("click_date")
    .parquet("/lake/clicks"))

The cost is 1 shuffle, plus a risk on busy days: a date with 10 times the usual traffic becomes 1 large file written by 1 task, and the stage waits for it. Adding .option("maxRecordsPerFile", 5000000) to the writer makes each task start a new file every 5 million rows, so a heavy date gets several bounded files.

Before choosing a column, look at how rows spread across its values, because each value's rows become 1 directory. Run this on store_sales_raw to see the rows each sale_date directory would hold, biggest day first, and compare the top of the list with the bottom.

Partitioning by a column with too many values

Every distinct value costs at least 1 directory and 1 file, however few rows it has. partitionBy("user_id") on a table with 2 million users creates 2 million directories. A read has to list 2 million paths before it scans a row, and the files average a few kilobytes, so opening each one takes longer than reading it.

This DataFrame has 50 partitions, and each partition holds rows from all 10 countries. How many data files does the write produce?

df.write.partitionBy("country").parquet("/lake/sales")

• Exactly 10, since Spark merges each country's rows into a single file
• Exactly 50, since each write task produces a single output file
• Exactly 60: 1 file per task plus 1 index file per country
• Exactly 500, since each of the 50 tasks writes into all 10 directories

To bring those 500 files down to 10, change which rows each task holds. repartition("country") hashes on the partition column, so all of a country's rows reach 1 task, and that task writes 1 file into its directory. The shuffle runs before the write, and explain() shows Exchange hashpartitioning(country, 200) above the scan. At most 10 of those 200 tasks hold rows, and the rest write nothing.

About this lesson

This section is part of the Partitioning: Intermediate lesson on DataDriven, the home of data engineering. Lessons cover SQL, Python, Spark, data modeling and pipeline architecture, and each section pairs its explanation with worked examples and practice you run as you go.